Dynamic output tensor shape

[sebastienwood]

Hi ! I'm writing a convolution-like operator using Stannum. It can be used throughout a neural network, meaning each layer may have a different input/output shape. When trying to register the output tensor, it leads to this error: AssertionError: Dim = -1 is not allowed when registering output tensors but only registering input tensors

Does it means I have to template and recompile the kernel for each layer of the neural network ?

For reference, here is the whole kernel/tube construction:

@ti.kernel
def op_taichi(gamma: ti.template(), mu: ti.template(), c: ti.template(), input: ti.template(), weight_shape_1: int, weight_shape_2: int, weight_shape_3:int):
    ti.block_local(c, mu, gamma)
    for bi in range(input.shape[0]):
        for c0 in range(input.shape[1]):
            for i0 in range(input.shape[2]):
                for j0 in range(input.shape[3]):
                    for i0p in range(input.shape[5]):
                        for j0p in range(input.shape[6]):
                            v = 0.
                            for ci in ti.static(range(weight_shape_1)):
                                for ii in ti.static(range(weight_shape_2)):
                                    for ji in ti.static(range(weight_shape_3)):
                                        v += (mu[bi, ci, i0+ii, j0+ji] * mu[bi, ci, i0p+ii, j0p+ji] + gamma[bi, ci, i0+ii, j0+ji, ci, i0p+ii, j0p+ji])
                            input[bi, c0, i0, j0, c0, i0p, j0p] += c[c0] * v
    return input


def conv2duf_taichi(input, gamma, mu, c, weight_shape):
    if c.dim() == 0:
        c = c.repeat(input.shape[1])
    global TUBE
    if TUBE is None:
        device = input.device # TODO dim alignment with -2, ...
        b = input.shape[0]
        tube = Tube(device) \
            .register_input_tensor((-1,)*7, input.dtype, "gamma", True) \
            .register_input_tensor((-1,)*4, input.dtype, "mu", True) \
            .register_input_tensor((-1,), input.dtype, "c", True) \
            .register_output_tensor((-1,)*7, input.dtype, "input", True) \
            .register_kernel(op_taichi, ["gamma", "mu", "c", "input"]) \
            .finish()
        TUBE = tube
    return TUBE(gamma, mu, c, input, weight_shape[1], weight_shape[2], weight_shape[3])

[ifsheldon]

If I understand correctly, you need to relate the shape of output with the shapes of inputs.

Tube manages Taichi field creation automatically on your behalf, so you need to tell it how to create these fields. -1 means an arbitrary dimension that is determined by the input shape, but it does not mean anything to the shapes of the outputs. Therefore, writing -1 dimension in input tensor shapes means "any dimension number that can be infer from the input tensor" but writing -1 dimension in output tensor shapes means "give me a tensor of an arbitrary shape".

You can use negative shapes other than -1, which means a dimension index.

For example, if one input tensor A has a shape (-2, ....), you can register output tensors with dimensions = -2. Say, the shape of one output tensor B is (10, -2, ...). This means the second dimension of B can be arbitrary but it is also determined by the shape of A. In this way, the shape of B has a correspondence with the shape of A, so Tube knows how to create Taichi fields for B.

[sebastienwood]

Hey thanks for the answer, this leads me to 2 questions:

• first, notice in the code snippet that I only add to the only registered "output" tensor (weirdly named "input") in the line before the return statement of op_taichi, the kernel. Does it mean I could not register "input" as an output tensor, but rather register it as an input tensor and thus not mind about shapes ?
• has I try to implement an operation similar to a convolution, I cannot express the output "input" using directly the input dimensions, i.e. I need to take into account the stride and other convolution parameters. Is it possible to express, e.g. with (10, f(-2), ...) where f(-2) allows me to compute on the fly the correct output shape ?  

[ifsheldon]

Does it mean I could not register "input" as an output tensor, but rather register it as an input tensor and thus not mind about shapes?

I think you don't have to return it. Why do you need to return the input field? but even if you absolutely need to return it, Tube currently does not support returning values yet.

has I try to implement an operation similar to a convolution, I cannot express the output "input" using directly the input dimensions, i.e....

Yeah, that's a problem. I've been thinking about a nice API since you filed this issue. If we allow shape calculation only for output tensors and intermediate fields, I guess a minor change in the code is sufficient. But if we also allow shape calculation for input tensors, designing a nice API gets a lot complicated.

[ifsheldon]

I wrote a prototype in the branch dim_calc. Maybe you can help test it. You'll need to write a subclass of DimensionCalculator, which needs only one method implemented. The arguments are dimensions (e.g., (-1, 10, -2, 20)) and shapes (concrete positive numbers of tensor shapes at run time, like (1, 2, 3, 6, 7)) of input tensors. And the output of it is dimensions (not shapes) so that it's more general.

I guess these arguments are sufficient and general enough to do any dimension calculation for output tensors. You can have additional attributes in a DimensionCalculator sub-class instance, like strides.

[sebastienwood]

I think you don't have to return it. Why do you need to return the input field? but even if you absolutely need to return it, Tube currently does not support returning values yet.

It's ok not to return it. The pipeline is as follow: I precompute some operation on the tensor input in Pytorch. The taichi kernel should add to input in place (input += ...). That's why I'm confused: it feels like I'm re-declaring input while trying to register it as the "output" tensor.

Unfortunately I cannot test the new branch at the moment, I'll let you know if I manage to !

[ifsheldon]

This is implemented in v0.9.0.