Passing "Any" as an argument to Relay ops

1

Hi, I have a use case where I need to pass Any to relay ops. Concretely, I want to make the max_loop_count below in the Torch frontend Any so that we can support truly dynamic for loop.

Right now max_loop_count is always const. But since in Python this is usually the size of some axis of tensors or the length of a list, which are dynamic, requiring the loop extent to a constant is severe limitation.

If I pass Any to relay.op.less above, I get

TVMError: Check failed: ObjectTypeChecker<TObjectRef>: :Check(ptr): Expect relay.Expr but get Any

How should I solve this problem? @zhiics @haichen @MarisaKirisame @kevinthesun @wweic

2

any is a dimension, it is not a value. I am confused - are you trying to pass a value as big as possible? That could be done by probably setting a very large number.

3

No. In my use case, Any results from tensor array stack op, which takes a list of tensors and stack them into a tensor one rank higher. Since the length of list is dynamic, the “stacked dimension” is Any. And then I want to loop over the stacked dimension.

Specifically, I want to support the following line in the Stacked LSTM example in the Torch repo.

The result of rnn_layer above is torch.stack, which in my translation to Relay gives a tensor with Any dimension. In the second iteration, the input tensor to the rnn_layer has Any as the size of the first axis, and I need to loop over this axis:

So even though Any in Relay is technically dimension, I think it makes sense to treat the dimension as value.

4

Does a shape_of function does the trick for you?

5

yes, I’ve just discovered shape_of function while looking at the test_any.py, and it worked!! Here is my new converter for aten::size(). If we discover Any as a result of infer_shape, I use shape_of function.

def _size(prelude):
    def _impl(inputs, input_types):
        shape = _infer_shape(inputs[0], prelude.mod)

        if any(map(lambda s: isinstance(s, tvm.tir.expr.Any), shape)):
            shape_dynamic = _op.shape_of(inputs[0])
            if len(inputs) > 1:
                axis = int(inputs[1])
                return _op.take(shape_dynamic, _expr.const(axis), 0)
            return shape_dynamic

        if len(inputs) > 1:
            axis = int(inputs[1])
            return shape[axis]
        return shape
    return _impl

I also got Stacked LSTM test working finally. Will send a PR soon.

1 Like
6

You probably know, but the more principled approach on LSTM is to use ADTs instead of dynamic shape tensor IMO. Doing so unlock optimization possibility such as loop fusion, dynamic batching and partial evaluation.

7

Yeah, right now my focus is to translate readily available Pytorch LSTM models. Python only has list and tuples :slight_smile: so I think I should stick with them for now. I’m also curious about “Tree LSTM” variant, although I don’t know what that is or how PyTorch people implements it.

8

think of it like this: a LSTM is a fold (recursive visit) on list, and a TreeLSTM is a fold (recursive visit) on rose tree.

data Rose a = MakeRose a (List (Rose a))

I had implemented treelstm a long time ago for relay benchmark, which is at here.

1 Like