Implement Symbolic Assumptions System

[jessegrabowski]

Description

Proposal

I would like to see a symbolic logic system for the mathematical properties of tensors implemented in pytensor. The name "assumptions" comes from the sympy system that does what I'm asking for. I guess it's actually a "type" system, but that name is quite overloaded.

Under this proposal, variable constructors would allow for a range of assumptions / types to be passed in. For example:

x = pt.tensor('x', shape=(None,), positive=True)

By assumption, values of x will always be positive. Other tags include non-negative, zero, non-zero, negative, non-positive. Under the sympy system, these values take one of three possible values: True, False or None, where None is interpreted as unknown or non-determinable.

The assumptions of a variable should be stored in an Assumptions Op, the same way the shape of an Op is stored in a Shape Op. Other Ops can then implement infer_assumptions, which tells us how they act on the various assumptions. Likewise, make_node will need to be expanded to include assumption logic.

The objective here is to have a symbolic "assumptions graph" that can be evaluated the same way the shape graph is evaluated. Calling x.assumptions should return a symbolic assumptions graph. Calling x.type.assumptions should return known static assumptions, or None where assumptions are not statically known.

Assumptions should have internal logic to maintain consistency. For example, a variable cannot be both positive and negative at the same time.

Elementwise Examples

If no assumptions are provided, nothing is returned.

# No assumptions given -- everything is unknown
x = pt.tensor('x', shape=(None,)) 
x.assumptions.eval() # {}

Some Ops can set assumptions even if assumptions of the inputs are unknown:

x = pt.tensor('x', shape=(None,)) 
z = pt.exp(z)
z.assumptions.eval() # {'positive': True}

Ops with multiple inputs should reason over the assumptions of both inputs:

x = pt.tensor('x', shape=(None,), positive=True)
y = pt.tensor('y', shape=(None,), positive=True)
z = x + y

# If x and y are both known to be positive, z can be assumed to be strictly positive
z.assumptions.eval() # {'positive': True}

Some Ops will destroy knowledge

x = pt.tensor('x', shape=(None,), positive=True)
y = pt.tensor('y', shape=(None,))
z = x + y

# y is unknown, so we can't assume that x + y is strictly positive
z.assumptions.eval() # {}

More Interesting Examples

The motivating examples for this feature relate to assumptions about matrices. For example, triangularity, positive-definite, negative-definite, positive semi-definite, banded, diagonal, topelitz, etc.

Users can define this manually

x = pt.tensor('x', shape=(None, None), positive=True, semidefinite=True)

Or it can be the result of an op.

x = pt.tensor('x', shape=(None, None))
L = pt.linalg.cholesky(x)
L.assumptions.eval() # {'lower_triangular: True}

# Question: should/can we propagate information back *up* the graph? Although x is not tagged, we can assume its PSD because it went into cholesky
x.assumptions.eval() # {'positive': True, 'semidefinite':True} ??

Multi-input Ops reason about the output assumptions

x = pt.tensor('x', shape=(None, ))
x = pt.diagonal(x) 
x.assumptions.eval() # {'diagonal': True}

y = pt.tensor('y', shape=(None, None))
z = x * y
z.assumptions.eval() # {'diagonal': True}

Other possible features

We could imagine using the assumptions to track the domain of a variable. This could help in rewrites involving conditions.

x = pt.random.beta(alpha=1, beta=1)
x.assumptions.eval() # {'domain': (0, 1), 'left_inclusive':False, 'right_inclusive': False}
z = pt.maximum(x, 0.5)
z.assumptions.eval() # {'domain': (0.5, 1), 'left_inclusive': True, 'right_inclusive': False}

Open Questions

• How should we handle matrix assumptions on higher order tensors? For example, what is the meaning of x = pt.tensor(shape=(None, None, None), positive=True, semidefinite=True). The first dim is always a batch dim? Or?
• Should this be a COp (like shape)? Or just do it in python?
• Should the assumption system duplicate type information? For example, should we have real and imaginary assumptions, or should assumptions also look at the dtype separately?
• How complex should the reasoning in the assumption system be? On every update, sympy updates every flag to be internally consistent. So when positive becomes True, nonnegative also becomes True, and negative becomes False.

Finally, the sympy team experimented with "new assumptions", presumably because there were flaws in the old assumptions. So copying their setup might not be the right way to go. I'm curious how this is handled in e.g. Maple and Mathmetica.

[ricardoV94]

Calling x.assumptions should return a symbolic assumptions graph

That's not a well defined operation unlike shape. If you evaluate the input of an Assumption Op you get an array, how do you extract assumptions at runtime? in other words what do you get from?

pt.assumptions(pt.abs(pt.matrix("x"))).eval()?

Shape is always well defined if you can only look at the input array at runtime.

We can't add arbitrary attributes to runtimes arrays because we want to remain compatible with the other backends. Even if we could it would be expensive to compute this at runtime (whereas shape isn't, which is why a ShapeOp is fine).

Also do you enforce the assumption at runtime (like specify_shape) or do you take it as a hint only?

Your proposal should contrast to the hint operation that existed in theano (or they considered having, don't recall now) but that was always a no-op that is only used during graph rewrites

Also making make_node more complex for all Ops seems like a bad strategy. I would write the logic to propagate assumptions in rewrites only. This allows us not doing the work when we don't want to.

[jessegrabowski]

That's not a well defined operation unlike shape. If you evaluate the input of an Assumption Op you get an array, how do you extract assumptions at runtime? in other words what do you get from?

pt.assumptions(pt.abs(pt.matrix("x"))).eval()?

At least {'positive':True}, plus also maybe implicated assumptions that are no longer None, like {'positive': True, 'negative': False}

[jessegrabowski]

Also making make_node more complex for all Ops seems like a bad strategy. I would write the logic to propagate assumptions in rewrites only. This allows us not doing the work when we don't want to.

I'm proposing a major extension to our API. We should always want to do this work. Doing the work in make_node makes the most sense, since the Node is the place where we reason about the symbolic relationships between inputs and outputs.

We could also do a multiple dispatch thing for each Op that performs the reasoning. But that seems more magical to me.

[ricardoV94]

https://github.com/Theano/Theano/blob/7505f23edd66fcdbbe04953f53c650b4d7cec4dd/theano/sandbox/linalg/ops.py#L34

[bwengals]

That's interesting, I searched theano for uses of Hint and didn't really see any. I could see users doing something like (GP example of course)

x1 = pt.dvector('x1')
x2 = pt.dvector('x2')
ell = pt.scalar('ell')

cov_func = pt.exp(-(x1[:, None] - x2[None, :])**2 / (2 * ell**2))
cov_func = pt.Hint(cov_func, symmetric=True, psd=True, postiive=True)

I can't imagine a way for PyTensor to be able to infer that cov_func will have those properties, unlike the outputs of pt.linalg.cholesky or pt.diag. Maybe pt.Hint could be applied internally to the outputs of those functions to accomplish the same behavior? Could a Hint op apply or remove tags, or attributes from @jessegrabowski s example?

[jessegrabowski]

pt.Hint applies tags, yes. I would prefer something like pt.specify_assumptions, since that matches the pt.specify_shape api. So in the simplest case, you could do exactly what you proposed. It's basically the same as the tags we currently (sometimes) use.

Ideally though, pytensor would know the following facts:

1. x[:, None] - x[None, ] -> sub.outer(x, y)
2. The output of sub.outer(x, x) or add.outer(x,x) (or just outer(x,x) for that matter) is symmetric
3. exp(sub.outer(x,x) ** 2) is psd
4. exp and square preserve symmetry

If it knew this, the cov_func could be automatically marked as positive semi-definite.

Fact (3) is pretty complex, probably we would need a rewrite to a specialized Op (GaussianKernel or something) that would do the marking.

[ricardoV94]

Fact (3) is pretty complex, probably we would need a rewrite to a specialized Op (GaussianKernel or something) that would do the marking.

I think you're missing in your mental picture the idea of rewrite Features and how they can be used to reason about the graph without having to introduce explicit constructs in the graph. I feel this will allow us to better separate concerns.

For instance ShapeFeature is pretty dumb right now, but one thing it could trivially do is create a list of all the sets of shapes that must be equal for a graph to be valid. Say you have pt.add(x, y), and pt.add(y, z) without any broadcasting. By transition, we know x.shape == y.shape == z.shape. That's easy to represent in a Feature like ShapeFeature, but not easy to represent in the graph directly.

[jessegrabowski]

Are there docs for Features, I didn't even know they exist

[ricardoV94]

I don't think there are docs, or good ones. They are objects attached to the FunctionGraph that can be triggered on specific FunctionGraph events via callbacks: A node is imported, a node is replaced, a node is removed..., etc.

ShapeFeature uses this to propagate shape graph across replaced nodes. a_shape = a.infer_shape(...) -> replace a -> b -> b_shape = a_shape), but more funky stuff like recording the changes so they can be undone later (we use this for the InteractiveRewrite).

The point is they are arbitrary classes, that exist on the FunctionGraph, not on the graph, and can reason and be told about changes in the graph as it gets rewritten.

[maresb]

I think it's important to temper ambitions on automatic inference of assumptions. Math is hard, Gödel, and blah blah, so it will never be as good as you might hope, and there isn't any "right way". Realistic expectations are basically that it will do well for whatever you hard-code, then it may randomly do a few clever things on top of that, and then it will rapidly taper off at some horizon (quite similarly to an LLM actually).

I'm curious how this is handled in e.g. Maple and Mathmetica.

Variables are by default assumed to be complex, so x² is not non-negative unless you declare x to be real. Occasionally Mathematica can infer something useful, but only in some low ratio of cases, so in the end I would usually never bother.

[jessegrabowski]

Yeah fair enough. I'm not interested in re-living the life of Bertrand Russel (for the sake of my marriage if nothing else). But I think we're doing a lot of checking in the linalg library that we could easily be automatically tracking. I think simple bookkeeping of properties like PSD and triangularity is well within reach. I also have in mind things like tracking whether a computation is linear or convex, like how cvxpy does, allowing us to dispatch to more solvers when you call minimize.

So there's low hanging fruit all over here, without getting into trying to build perpetual motion machines.