Dynamic shape handling

[zezhang]

Is there any guidelines to follow for handling the dynamic shape (e.g. reshaping a dynamic tensor)?

In addition to using the shapeOp + sliceOp, is there any other way to get a specific dim size?

Is there a good way to construct 1-D tensor for the ShuffleLayer's dynamic shape input? Currently, I'm using shapeOp + sliceOp + concatOp to construct 1D tensor for mixing dynamic batch with static dims, but it's quite complicated. It would be great if there is an easier way to construct 1D tensor, and a direct way get the size of a particular dim from a tensor.

[pranavm-nvidia]

You should be able to do dynamic reshaping by manipulating the tensor.shape as though it were a list. For example, you can flatten a tensor by taking the product of its shape:

>>> import math
>>> import tripy as tp
>>>
>>> def flatten_dynamic(x):
...     return tp.reshape(x, [math.prod(x.shape)])
...
>>> x = tp.ones((1, 2), dtype=tp.float32)
>>> flatten_dynamic(x)
tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=[2])

Dynamic shapes are really only relevant in compiled mode since eager mode always uses dynamic shapes. You can compile the above function for a range of input shapes and it should still work as expected:

>>> compiler = tp.Compiler(flatten_dynamic)
>>> compiled_flatten = compiler.compile(tp.InputInfo(shape=[(1, 2, 3), 2], dtype=tp.float32))
>>>
>>> x = tp.ones((1, 2), dtype=tp.float32)
>>> compiled_flatten(x)
tensor([1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=[2])
>>>
>>> x = tp.ones((3, 2), dtype=tp.float32)
>>> compiled_flatten(x)
tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000], dtype=float32, loc=gpu:0, shape=[6])