Automatic operation fusion

[Bihaqo]

For example, instead of writing a separate function project_matmul we can implement something like project_matmul = fuse(lambda a, b: project(a, b), lambda b, c: matmul(b, c))

To do that lets express each operation as a sequence of recurrent steps of the form

def recurrent(a, b):
  res = 1.0
  for a_core, b_core in zip(a.tt_cores, b.tt_cores):
    res = einsum('rule', a_core, b_core, res)

and of independent steps

def independent(a, b):
  res_cores = []
  for a_core, b_core in zip(a.tt_cores, b.tt_cores):
    res.append(einsum('rule', a_core, b_core))

Then, we can automatically concat einsum-s of individual operations into a single big einsum (per core), and by using opt_einsum guarantee that the restulting einsum will be fast.

From the top of my head we can support any combinations of 1) matmul(A, B) 2) add a + b 3) elementwise product a * b

Additionally, as the last operation of the combination, we can support 1) dot product a^t b 2) gram matrix G_ij = ai^t bj 3) projection on the tangent space P_x y 4) trace

By combining this ops we can for example automatically get fast versions of 1) x^t A y (already implemented as a separate fast operation) 2) ||A B|| 3) A B x 4) P_x A y (already implemented) 5) ||a b|| 6) Px A B y 7) ||A + B|| 8) P_x (a b) 9) x^t A B y 10) ||(Ax) * (By)|| 11) trace(A^T B A)

Does anyone need this?

[Bihaqo]

A potential way to design the API:

with t3f.Fuse() as f:
  Ax = t3f.matmul(A, x)
  xAx = t3f.flat_inner(x, Ax)
  fast_xAx = f.optimize(xAx)