Dimension information in TensorType

[aseyboldt]

I don't know if this was discussed before, but I think it could be really interesting and useful to store more information about dimensions in the TensorType.

We could for instance define objects for individual dimensions roughly like this:

class Dimension:
    name: Optional[str]
    length: Optional[str]

We could then allow dimension information in the TensorType directly.

country = Dimension(name="country")
x = at.dscalar(name="x", dim=country)

Which would give us a TensorType(np.float64, shape=(None,), dims=(country,))

Many ops could then figure out the dimensions for output types. at.zeros_like for instance could tell us that the dimension of the output is the same as the dimension of the input, rather than just an unknown shape with the same rank as the input.

# This could then know 
y = at.zeros_like(x)
y.type.dims == (country,)

Just for debugging alone I think this would be really helpful. If there are some shape mismatches the error message could tell us which dimensions are the problem.
It would also allow pymc to infer dims for eg deterministics.
We could add at.dim_add(x, y) or so that does broadcasting based on those dimensions instead based on the position of the dimensions, which I think can prevent a whole class of common bugs.
I can also imaging that this could be useful for optimizations (especially reusing memory from old computations). If we have dimensions, we can tell that two variables have the same shape, whatever that shape might be.

[ricardoV94]

Not sure if you followed the recent changes in static shape types: https://aesara.readthedocs.io/en/latest/extending/type.html

In general we are trying to make the output types more clever about static shapes, so the following should work like this:

import aesara.tensor as at

x = at.tensor(dtype="float64", shape=(50,), name="x")
y = at.zeros_like(x)
y.type  # TensorType(float64, (50,))  <- Not yet there

This does not require any new functionality, only improving the output of made_node. This PR, for instance, tries to do this for the Elemwise operators: #821

Your suggestion goes one step further and gives names to these dimensions, but I feel part of it was related to being clever about the propagation of static type shape, which we are already working on.

[ricardoV94]

For shape info across unrelated Ops, one can also use SpecifyShape. What we do with that in the compilation machinery is a different matter though.

country_dim = at.iscalar("country_dim")  # or a constant, or a shared variable
x = at.specify_shape(at.vector("x"), shape=(country_dim,)
y = at.exp(x)
...

[aseyboldt]

Ah, I hadn't seen this, thanks. Seems closely related.
Dimensions are different than static shapes though, for instance a type TensorType(np.float64, shape=(50,), dims=(country,)) would be different from a TensorType(np.float64, shape=(50,), dims=(state,)).

[brandonwillard]

Ah, I hadn't seen this, thanks. Seems closely related. Dimensions are different than static shapes though, for instance a type TensorType(np.float64, shape=(50,), dims=(country,)) would be different from a TensorType(np.float64, shape=(50,), dims=(state,)).

That can be modeled equally well by a distinct TensorType subclass, and it probably should be, so that we can make use of Python's type system (e.g. mypy). (If there are any issues with the use of a TensorType subclass, then those are bugs we need to address.)

Overally, we're trying to move away from the mostly redundant Type framework that we've inherited from Theano. See #134.

[brandonwillard]

See #1089 (comment) for an illustration and more information about the relationship between the dimension information mentioned here and the existing specify_shape and shape inference.