ENH: Implement COLA library rewrites for linear algebra functions

[jessegrabowski]

Description

This paper and this library describe and implement a number of linear algebra simplifications that we can implement as graph rewrites. This issue is both a tracker for implementing these rewrites, and a discussion for how to handle them.

Consider the following graph:

x = pt.dmatrix('x')
y = pt.diagonal(x)
z = pt.linalg.inv(y)

If we could promise at rewrite time that y is diagonal, we could re-write the last operation as z = 1/y, exploiting the structure of the diagonal matrix. Other non-trivial examples exist, for example:

x = pt.dmatrix('x')
y = pt.eye(3)
z = pt.kron(x, y)
z_inv = pt.linalg.inv(z)

If we could promise at rewrite time that z is block diagonal, we could rewrite z_inv = pt.kron(pt.linalg.inv(x), y), which is a much faster operation (since x is 3x smaller than z).

The linked paper and library list a huge number of such shortcut computations. The following is a list version of Table 1 from the paper. Under each function are the types of matrices for which a rewrite rule exists to accelerate the function. It would be nice to collaborative update this list with links to the COLA library where the relevant rewrite is implemented:

Thanks to @tanish1729 for compiling a list of links to relevant rewrites. Missing links indicate no direct implementation

• cholesky(A)

  • Identity
  • Diagonal
  • Kronecker Product
  • Block Diagonal

• inv(A)

  • Identity
  • Diagonal
  • Triangular
  • Permutation
  • Convolution
  • Sum
  • Product
  • Kronecker Product
  • Block Diagonal
  • Concatenation

• eigs(A)

  • Identity
  • Diagonal
  • Triangular
  • Convolution
  • Sum
  • Kronecker Product
  • Block Diagonal

• diag(A)/trace(A)

  • Identity
  • Diagonal
  • Sum
  • Scalar Product
  • Triangular
  • Permutation
  • Kronecker Product
  • Block Diagonal
  • Concatenation

• logdet(A)

  • Identity
  • Diagonal
  • Triangular
  • Permutation
  • Scalar Product
  • Kronecker Product
  • Block Diagonal

• Unary Operations (such as log, pow, exp, sqrt)

In addition, potential re-writes could also be applied to Topelitz and Circulant matrices, although these are not covered by COLA.

The wrinkle to all this is that we would need more information about matrices as they enter and exit Ops. Right now, we're using tags to accomplish rewrites like this, see for example #303 and #459 . Some of these rewrites might be possible to do via inference. For example, a pt.linalg.block_diag Op always returns a block diagonal matrix, as does pt.kron(pt.eye, A). pt.diagonal always returns a diagonal matrix, as does pt.eye; pt.linalg.cholesky always returns a triangular matrix, etc. Other potential types, like block, positive, definite, psd, topelitz, circulant, etc, would be less trivial to automatically detect.

The other issue is that as these type tags proliferate, we become more and more locked into a somewhat hack-y system for marking things. Perhaps putting some thought into how to handle this now will save some refactoring headaches down the road?

[ricardoV94]

I think the easiest short term path is to add a feature (like the shape feature) that
does/keeps track of this more specialized type inference.

For instance when you replace a variable by another you can assume what you inferred for the old variable holds for the new.

When a new variable is added you can compute new stuff. Like x * 5, where x is positive/diagonal is also positive/diagonal. If it's positive it's not negative and so on.

Some type info needs to be seeded by the user initially, as they won't emerge naturally from realistic graphs. Right now we use tags for that but maybe we could add a dummy pass through Op, like that is ultimately discarded but feeds initial type info. User could say x = pt.hint(x, "positive") or pt.hint(x, x>0).

Or we could introspect asserts like pt.assert(x, x>0) but we probably don't want to keep most of these asserts in the final graph.

[jessegrabowski]

This is how sympy handles it, with their assumptions system. For example:

import sympy as sp
x = sp.Symbol('x', positive=True)
x._assumptions0

Gives a list of implications of x being positive:

{'positive': True,
 'extended_nonnegative': True,
 'commutative': True,
 'nonpositive': False,
 'extended_positive': True,
 'nonnegative': True,
 'nonzero': True,
 'imaginary': False,
 'infinite': False,
 'extended_nonpositive': False,
 'real': True,
 'extended_negative': False,
 'complex': True,
 'extended_nonzero': True,
 'zero': False,
 'negative': False,
 'hermitian': True,
 'finite': True,
 'extended_real': True}

All of these have implications for their algebraic simplifications. For example, it will refuse to simplify (x ** 0.5) ** 2 to x if it doesn't know that x is strictly positive.

It's actually a nice system but I read they were trying to revamp it, so I'm curious what they think the big failures of it are. It might be nice to try to set up a call with their devs to talk about their lessons from going this route before we commit to it.

I also think something like this is really where egg rewrites can shine. Especially if we're just dealing with matrices, we shouldn't have to reason about shapes too much (although Blockwise somewhat ruins that...)

[ricardoV94]

It might be nice to try to set up a call with their devs to talk about their lessons from going this route before we commit to it.

Definitely

[tanish1729]

hey @jessegrabowski! are there any more related PRs that i could try working on for this project?

otherwise, from what i understand, this is mainly focused on graph rewriting? and we want to be using the computations suggested in the paper in pytensor

[jessegrabowski]

Yeah basically we want to introduce simplifications related to compositions of linear algebra operations via rewrites. The COLA paper suggests a big pile, so that's what I suggested as a "roadmap".

The only active one right now is #622 , you can look here for a rough template of a rewrite looks like. Maybe a "simple" rewrite to get your feet wet would be to change det(diag(x)) to prod(x). The rewrite should look for a det Op, then check if the parent is a diag op. If it finds that, it should return the product of the input to diag.

Obviously there are other ways we could look for a diagonal matrix to arise (for example, eye * x), but that would be a good place to get the ball rolling.