`Subtensor` shape incorrect shape inference

[larryshamalama]

There seems to be an issue with static shape inference for Subtensor ops. Below is an example provided by @brandonwillard in PR #901:

import aesara.tensor as at

y = at.tensor(dtype="float32", shape=(2,))
y.type.shape
# (2,)

# We know that the shape of this should be equal to the shape of `y`
z = y[slice(None)]

aesara.dprint(z, print_type=True)
# Subtensor{::} [id A] <TensorType(float32, (None,))> ''   
#  |<TensorType(float32, (2,))> [id B] <TensorType(float32, (2,))>

z.type.shape
# (None,)

See comment here.

[larryshamalama]

I was able to give this issue some thought and hopefully this leads to something.

Looking at the make_node method in Subtensor, I see that the shape argument to TensorType is a list of bools (see lines 731 to 734), which gets converted to 1s or 0s. However, it is my understanding that the list provided to the shape argument can also contain ints when a dimension has a fixed shape that is not 1. It seems that fixed shapes that are not 1 (such as 2 in the example above) becomes None because it is assigned broadcastable=False.

Is there a way to include static shape information in the broadcasting list through the for loop? Now we have something like:

for i, (p, bc) in enumerate(zip(padded, x.type.broadcastable)):
    if bc:
        ...
        broadcastable.append(True)
    else:
        broadcastable.append(False)

Would something along the following lines be a good blueprint?

for i, (p, bc, shp) in enumerate(zip(padded, x.type.broadcastable, x.type.shape)):
    if bc:
        ...
        broadcastable.append(True)
    else:
        if shp is not None: # static shape that's not broadcastable (> 1?)
            # somehow figure out new shape using old shape `shp` and slice `p`...
            ...
            broadcastable.append(new_shp)
        else:
            broadcastable.append(False)

If these thoughts are correct, my first guess is that figuring it out new_shp will be somewhat challenging

[brandonwillard]

Looking at the make_node method in Subtensor, I see that the shape argument to TensorType is a list of bools (see lines 731 to 734), which gets converted to 1s or 0s.

That alone is a bug; the broadcastable variable should be renamed to shape/static_shape and Nones used in place of Falses/1s used in place of Trues. (It could also be changed to tensor(dtype=x.type.dtype, broadcastable=broadcastable), but the broadcastable keyword argument is being deprecated.)

[brandonwillard]

If these thoughts are correct, my first guess is that figuring it out new_shp will be somewhat challenging

Subtensor.infer_shape has all the necessary shape-determining logic in it, so you can always look to that for answers about the basic logic; however, that method is designed to work with symbolic values, and the Type.shape calculations need to be performed on ints and Nones instead.

As we've done in other cases, Op.infer_shape can be used explicitly to determine the static shape values of an Op's outputs, but we should avoid that as much as possible right now, because it involves Function compilation when Ops are created.

[larryshamalama]

Please correct me if I am misunderstanding.

Subtensor.infer_shape has all the necessary shape-determining logic in it, so you can always look to that for answers about the basic logic; however, that method is designed to work with symbolic values, and the Type.shape calculations need to be performed on ints and Nones instead.

As we've done in other cases, Op.infer_shape can be used explicitly to determine the static shape values of an Op's outputs, but we should avoid that as much as possible right now, because it involves Function compilation when Ops are created.

My current understanding from your comment is that I should borrow the idea behind shape inference from Subtensor.infer_shape rather than use it explicitly because the method works with symbolic shapes, which static shapes are not. Is this correct?

[brandonwillard]

Please correct me if I am misunderstanding.

Subtensor.infer_shape has all the necessary shape-determining logic in it, so you can always look to that for answers about the basic logic; however, that method is designed to work with symbolic values, and the Type.shape calculations need to be performed on ints and Nones instead.
As we've done in other cases, Op.infer_shape can be used explicitly to determine the static shape values of an Op's outputs, but we should avoid that as much as possible right now, because it involves Function compilation when Ops are created.

My current understanding from your comment is that I should borrow the idea behind shape inference from Subtensor.infer_shape rather than use it explicitly because the method works with symbolic shapes, which static shapes are not. Is this correct?

Yes, exactly.

[ricardoV94]

We should open a general issue for all the Op's that don't propagate static shape information (and decide which ones, if not all, should). Subtensor is only one such case.