Static shapes

[covertg]

Description

From this discussion on pymc discourse, "static shapes are quite recent and we are still updating Ops to reflect them when possible." -@ricardoV94

Is help still needed with ensuring static shape output from Ops where possible? If so - here's a quick (obviously incomplete!) start on listing out ops to work on.

No static shape output (todo):

• Eye
• conv ops (Output shapes for conv do not pass-by static dimensions #150)

No static shape output (and perhaps to remain as such):

• Subtensor (indexation)

Could fail earlier:

• Dot: outputs correct static shape if the input shapes align, but outputs incorrect static shape if input shapes are wrong. Error at eval time
• SpecifyShape: similar to Dot. ptt.specify_shape(np.ones((2,3)), shape=(3,3)).type # tensortype shape (2,3)

Recently addressed:

• Join Compute static shape types in outputs of Join #164
• MaxAndArgmax Propagate static shape in MaxAndArgmax #206

Hopefully this is the right place to start the discussion. If help is wanted, I would appreciate some guidance on how to go about this in a more principled way -- where/how to look, what a fix might look like, etc -- and of course any other suggestion on Ops that don't yet support static shape output.

[ricardoV94]

Hi @covertg, thanks for the write-up.

From the listed missing cases note that Subtensor is particularly challenging. There are many forms of indexing, from slicing to advanced integer and boolean indexing, as well as mixtures of these (just like numpy)

This makes it challenging to infer output static shape without actually performing the subtensor operation. In addition this would only be possible when the respective inputs are constant, and even accessing that may be too costly. Given that, I think we will likely keep undetermined output static shapes.

The Eye case should be easier to support just like we do with Alloc, by trying to constant fold the input n

MaxAndArgmax should be straightforward. We already propagate broadcastable info, just need to propagate static shape instead.

---

About the "bugs" those are not really wrong. If we know the inputs are incompatible with the output static shape we can raise early, but it's not otherwise a bug. PRs to raise early are welcome though.

[covertg]

Ah, thanks for all the clarification!
I hadn't appreciated all the complexity in Subtensor. That does seem rather difficult.
I've updated the original post to represent the subtensor question, and to reframe the "bugs" not bugs.

I'm interested in trying to help with this, so I'll have an initial PR for MaxAndArgmax soon!

[covertg]

I've also noticed that SharedVariable does not store a static shape when it's provided a numpy array. I'm not as familiar with how shared variables are used -- this came up because one of the maxandargmax tests used shared() rather than as_tensor_variable().
Do we want to store static shape for instances such as

>>> pt.shared(np.array([1,2,3]))
<TensorType(int64, (?,))>

?

[ricardoV94]

Shared variables can be resized by default so they can't have static type shapes. Users can however define a static shape when creating it

pytensor.shared(np.ones(10), shape=(10,))

[ricardoV94]

Having said that Subtensor cases can be too hard, there are some simple ones that should be straightforward. For instance when doing a simple scalar indexing, we know the other dims remain the same:

import pytensor
import pytensor.tensor as pt

idx = pt.scalar("idx", dtype="int64")
x = pt.tensor("x", shape=(5, 4, 3))
y = x[idx, :, :]
print(y.type.shape)  # (None, None)  -> Could be (4, 3)

y = x[:, idx, :]
print(y.type.shape)  # (None, None)  -> Could be (5, 3)

y = x[:, :, idx]
print(y.type.shape)  # (None, None) -> Could be (5, 4)

We can also support slicing with constant scalars

y = x[1:, :, :]
print(y.type.shape)  # (None, None, None) -> Could be (4, 4, 3) or at least (None, 4, 3)