Enable complete static shape information at the type level

[brandonwillard]

This PR addresses #431 and—as a result—#608 by adding full static shape information at the Type level.

I recently noticed that something like this was attempted a while back, and that some of the relevant issues are described in typeattr.txt. As mentioned in that document, the main issue is that a.type == b.type is not suitable for use when Types have varying degrees of type/shape information.

This varying degree of type information is currently dealt with by making all Types conform in the case of replacements. Types can differ during rewrites when, for instance, the replacement Variable is the output of an Op that's better at inferring TensorType.broadcastable information (i.e. static shape information). In these cases, if a rewrite/optimization replaces a Variable of Type TensorType("float64", (True, False)) with a Variable of Type TensorType("float64", (False, False)), it will apply a Rebroadcast Op in order to "broaden" the original (True, False) broadcast pattern (i.e. shape equal to (1, s) for some s) so that the replacement Type is equal to the replaced variable's original Type.

There are at least two things wrong with this approach: 1) Rebroadcast is basically just a redundant Assert Op, and 2) Type "broadening" isn't particularly productive or useful. We should only be "narrowing" the Types. In other words, if we get more information about a Type, we should use it. This puts less pressure on each individual Op to operate at an equally precise level of shape inference, and it more explicitly draws out contradictions/inconsistencies between Op implementations.

To summarize the above, we need to "narrow" Types as we refine a graph. We're talking about a simple form of type inference, essentially.

In order to do this, an explicit "information"-based ordering over Types is required. For example, a.type < b.type is true when a is a more specific type than b, and, as a result, a can be replaced by b under the assumption that both types represent the same term. Since this is a high-level assumption we necessarily make when performing replacements, we're basically updating our type information about terms when we encounter these cases.

For example, t1 = TensorType("float64", [False, False]) can be replaced by t2 = TensorType("float64", [True, False]), because t2 has a more specific broadcast pattern than t1. In other words, t1 > t2 according to the above.

The implementation in this PR uses a new Type.is_super method to implement the < ordering in the above example, but it can easily be replaced by a more convenient __lt__ implementation. (NB: the name of the method is liable to change.)

After making these important Type system changes, this PR introduces an entirely new Type to implement fixed-shape types; however, since the existing TensorType already has partial shape information (i.e. TensorType.broadcastable), it might be better to extend that by adding a TensorType.shape attribute and making TensorType.broadcastable a property derived therefrom. This aligns closer with the approach described in the old typeattr.txt file above, since the TensorType.shape would need to encode missing shape information with something like Nones (e.g. TensorType.shape = (None, 1, None) would be equivalent to TensorType.broadcastable = (False, True, False)).

• Finish refactoring code that uses strict Type equality
• Add tests for new fixed-shape Types
• Consider making all TensorTypes with all(TensorType.broadcastable) is True fixed-shape tensors, since their shapes are fixed and known already.
  This would involve overriding TensorType.__new__ so that it returns the appropriate types.
• Clarify/fix the Type.filter_variable/Type.convert_variable interface. Currently, these methods overlap a lot and their exact purposes aren't clear and their usages seem somewhat inconsistent with regard to types and the like.
• Consider removing distinct FixedShapeTensorType and extending the existing TensorType to include a shape field.

[codecov]

Codecov Report

Merging #711 (f266c52) into main (51792fe) will increase coverage by 0.18%.
The diff coverage is 87.14%.

@@            Coverage Diff             @@
##             main     #711      +/-   ##
==========================================
+ Coverage   78.17%   78.35%   +0.18%     
==========================================
  Files         152      152              
  Lines       47663    47679      +16     
  Branches    10881    10879       -2     
==========================================
+ Hits        37260    37360     +100     
+ Misses       7846     7772      -74     
+ Partials     2557     2547      -10     

Impacted Files	Coverage Δ
aesara/d3viz/d3viz.py	23.68% <0.00%> (ø)
aesara/graph/features.py	65.95% <ø> (ø)
aesara/link/basic.py	85.21% <0.00%> (ø)
aesara/link/vm.py	87.19% <ø> (ø)
aesara/tensor/nnet/batchnorm.py	77.04% <ø> (ø)
aesara/tensor/nnet/neighbours.py	91.26% <0.00%> (ø)
aesara/tensor/nnet/opt.py	42.96% <0.00%> (ø)
aesara/tensor/math_opt.py	86.23% <22.22%> (-0.01%)	⬇️
aesara/compile/function/types.py	78.62% <33.33%> (ø)
aesara/scan/opt.py	80.31% <33.33%> (ø)
... and 57 more

[ricardoV94]

This is really interesting progress.

I don't have anything to add just yet, but adding this link to an old Theano devs discussion that seems relevant after this PR (when mixing fixed and non-fixed shape variables): https://groups.google.com/g/theano-dev/c/OKiVysM4ySg/m/6Q4iHLseBwAJ

[kc611]

Overall this PR (is extremely huge and) a really interesting change.

[kc611]

This aligns closes with the approach described in the old typeattr.txt file above, since the TensorType.shape would need to encode missing shape information with something like Nones (e.g. TensorType.shape = (None, 1, None) would be equivalent to TensorType.broadcastable = (False, True, False)).

I have a question over here. If we have the full static shape information at the Type level, what are the cases in which we might be needing this ?

I might be misunderstanding what we mean by full static shape information, I assume it is the literal information about the shape. (as a tuple of integers or something which helps us derive that.)

[brandonwillard]

This aligns closes with the approach described in the old typeattr.txt file above, since the TensorType.shape would need to encode missing shape information with something like Nones (e.g. TensorType.shape = (None, 1, None) would be equivalent to TensorType.broadcastable = (False, True, False)).

I have a question over here. If we have the full static shape information at the Type level, what are the cases in which we might be needing this ?

I might be misunderstanding what we mean by full static shape information, I assume it is the literal information about the shape. (as a tuple of integers or something which helps us derive that.)

Yes, that's exactly what it is; one can create a TensorVariable (or a subclass thereof) that has fixed, concrete shape information, such that var.shape would return a tuple of ints.

Aside from providing exact shapes upfront, it can also simplify a lot of the shape inference and remove the need for some graph elements (e.g. Shape, Rebroadcast, and similar Ops).

More immediately, we simplify our Op implementations, which currently have no consistent means of modeling fixed-size/shape terms. For instance, we're forced to use vararg inputs, or Op-level constants that specify the number of inputs/dimensions (e.g. Op.ndim), in order to model shape-like arguments, but those approaches quickly become cumbersome or even not applicable in the face of multiple such inputs. Join, Split, MakeVector, and Alloc are all good examples. With these changes, it would be possible to use a single fixed-length/shape input in these Ops (e.g. MakeVector's *inputs could be input, if input were fixed-length). (Note that this could also be addressed with support for a tuple-like type.)

For the same reasons, to fix #608 we would need to change RandomVariable so that it carries around the concrete length of its size input, just like those other Ops; however, we can address this—and much more—with the changes here.

There are also a few other related interfaces (e.g. #93) that can be improved along the way, like TensorConstant.shape, which can be made consistent with its premise of modeling a known constant and—thus—a known shape.

[twiecki]

Nice, this is a big one 🥳 .

[twiecki]

I know you don't like to talk about JAX ;) but does that help at all with omnistaging?

[brandonwillard]

I know you don't like to talk about JAX ;) but does that help at all with omnistaging?

It actually could, but there are still a few follow-ups needed in order to consistently propagate this new static shape information. I'll try to provide those in the next couple days.