Allow specification of !=1 Type shape

[ricardoV94]

Brought up in #1089

Knowing that two tensors will not broadcast (i.e., dimensions != 1) allows for more succinct graphs in some cases as well assumptions in rewrite. In fact many old rewrites developed in Theano relied on something like that (although it was a bit murky, see #1089 for more details), and are nowadays incorrect. Here is one known example in the codebase:

aesara/aesara/tensor/basic_opt.py

Lines 1695 to 1705 in 7f8af9b

	if values_ref.type.broadcastable == out_type.broadcastable:
	# The assumption here is that `values_ref` already has the same shape
	# as `shape_ref`, so a `fill`/`Alloc` is unnecessary.
	# XXX FIXME TODO: The only way this can be determined is if one
	# absolutely knows that the shapes of `shape_ref` and `values_ref` are
	# equal.
	# This is an old rewrite, and it's only a
	# "specialization/stabilization", so we're going to leave it be for
	# now.
	return [values_ref]

A common pattern, is that when you have an Elemwise operation between two inputs, and they have the same broadcastable pattern, you can assume they have the same shape (otherwise it would be an invalid graph). With current Aesara Types you can not make this assumption because there is no way to specify shape!=1 at the Type level, without specifying a specific shape.

[ricardoV94]

This touches on a more broad question of what type of information we want to carry at the Type level. Other examples include value domains (e.g., strictly positive, unit domain, semi-definite matrix, and so on), which when known allow for specific rewrites / simplification that could otherwise not be performed.

E.g., #942

[rlouf]

+1 to include value domains. It would potentially allow us to remove distribution-specific logic for the transforms in aeppl, and statically check that the graph is compatible with each variable's parameter constraints and distribution support.

[brandonwillard]

This touches on a more broad question of what type of information we want to carry at the Type level.

Remember, there is direct overlap between Type-level information (e.g. TensorType.shape) and "graph"-level information (e.g. an Assert with constraints encoded by graphs themselves). For instance, both types of information are available and used at compile-time, and one can generally be converted to the other format.

Since we have the canonicalization and constant-folding mechanics implemented at the graph-level, it's possible for us to reuse and automate a lot more in that context; otherwise, we don't want to take the idea of Type-level information too far and end up reimplementing those things by—for example—coming up with new ways to represent x > 1 at the Type-level when we already have at.gt(x, 1) at the graph-level.

[ricardoV94]

Do we have any examples of using asserts like that? I imagine constant-folding could mess up things if it removes asserts too early, but maybe that does not affect the graphs where this would matter.

[rlouf]

An intermediate solution would be to represent the constraints as classes:

class interval():
    def __init__(lower, upper):
        self.upper == upper
        self.lower == lower

    def __call__(x: TensorVariable):
        return at.ge(x, self.lower) & at.le(x, self.upper)

unit_interval = interval(0, 1)

x = at.scalar()
x = CheckConstraint(x, unit_interval)

This way we can both reason on the graph (by adding assert statements where needed) and in terms of types.

[brandonwillard]

Do we have any examples of using asserts like that?

We don't have anything that reasons about sets of Assert constraints in a graph. That's what we need to add. I'm saying that the exact same functionality is what we need for shape constraints, so we probably shouldn't consider adding anything new/different to specifically encode and/or process shape-related constraints.

I imagine constant-folding could mess up things if it removes asserts too early, but maybe that does not affect the graphs where this would matter.

Currently, constant folding can—and should—remove Asserts when their conditions are computable at compile-time. What would be the point of evaluating things like that at run-time?

Again, I'm talking about the approach mentioned in #942 (reply in thread) and further detailed in #1089 (comment). More specifically:

use Assert (or a subclass) to represent constraints like X.shape[i] != 1 with equivalent in-graph Ops like at.neq(X.shape[i], 1). Next, add a Feature that picks up these nodes and adds an entry mapping X.shape[i] to (at.neq, 1) in a dict. There's your constraint store. Finally, create a local rewrite that tracks all other boolean Ops and replaces matching conditions in the constraint store with conforming True values. Now, those entries can be constant folded–or explicitly removed by the rewrite itself in some cases.

[brandonwillard]

Knowing that two tensors will not broadcast (i.e., dimensions > 1) allows for more succinct graphs in some cases as well assumptions in rewrite. In fact many old rewrites developed in Theano relied on something like that (although it was a bit murky, see #1089 for more details), and are nowadays incorrect.

The last part of this is not true, especially for the example in #1089.

Theano had exactly the same underlying issues with compile-time broadcasting that we currently have, and it's because we're still using Theano's old broadcasting assumptions: i.e. we interpret TensorType.shape[d] == None as False in TensorType.broadcastable, so TensorType.shape and TensorType.broadcastable represent the exact same information to any logic based only on the latter. You even created a discussion about this fact: #1170.

The best one could say is that we want a consistent interpretation of TensorType.shape[d] == None as "no shape information at all", but we obviously don't have that.

What we actually did was allow Elemwise.perform (and its C equivalent), and a few other Ops, to work without static broadcast information (i.e. without needing to know that a dimension is not equal to one), and that only changed the way in which these old Theano errors appear. See #1089 (comment) for an example of the same incorrect broadcasting assumptions in Theano, or just take a look through Theano's issue tracker.

With current Aesara Types you can not make this assumption because there is no way to specify shape>1 at the Type level, without specifying a specific shape.

This is also a meaningless statement, because we are making the same broadcasting assumptions as Theano, and especially in the only places where it matters (e.g. the old rewrites and constructors that need broadcasting information).

Why mention all this? We really need to understand why we want to add constraints like shape > 1 or shape != 1 and what we can get from them, and so far this discussion has only offered some awkward misconceptions.

Here's one real reason for adding constraints: we want to correct the interpretation error mentioned above (i.e. the subject of #1170) and distinguish between unknown shapes (i.e. None values) and non-broadcastable shapes. Doing so would allow us to raise better errors at compile-time, avoid some rewrite inconsistencies, and better regulate the use of our new "dynamic broadcasting" support and the undesirable graph additions it may entail.