cooijmanstim / einig Goto Github PK

Signatures like bij,bjk,bkl->bil where some axis appears in more than two inputs trigger an "exponential-space" implementation that broadcasts everything into an intermediate tensor bij,bjk,bkl->bijkl and then reduces it bijkl->bil. This should be pretty easy to detect and break up into explicit steps, e.g. bij,bjk->bik followed by bik,bkl->bil.

einsum("bij,(ij)k->bk", x, y) where the leading axis of y is actually in the same space as the trailing axes of x, e.g. x has shape [B, I, J] and y is the result of flatting the leading axes of something shaped [I, J, K].

It is well known that tf.gradients (and its relatives in other frameworks) is an inflexible beast. Just how inflexible doesn't become clear until a better way is shown. Here's an einsum-inspired notation for multiplication by derivatives.

tf.gradients(y, x, grad_ys=z) computes something like einsum("i,ij->j", z, dy/dx) without the explicit computation of dy/dx. This is the only case supported.

Suppose we write this eingrad("i,i%j->j", z, y, x), where eingrad will take care of calling out to tf.gradients appropriately. The % symbol separates operands like , does, but additionally indicates that the operand to the left of it should be differentiated with respect to the operand to the right of it.

Full jacobians for vector y: eingrad("i%j->ij", y, x)
Hessians for scalar y: eingrad("j%k->jk", eingrad(",%j->j", y, x), x)
Batched gradients for scalar y: eingrad("b%bj->bj", y, x)
Jacobian-vector product: eingrad("i%j,j->i", y, x)
Batched Jacobian-vector product: eingrad("bi%bj,bj->bi", y, x)

These cases are ridiculously difficult and error-prone to implement with bare tf.gradients, requiring various kinds of loops, transpositions and reshapes and acute awareness of the fragile notion of connectedness in symbolic graphs (e.g. tf.gradients(y, tf.identity(x)) or tf.gradients(y, x[0]) does not work at all, whereas tf.gradients(tf.identity(y), x) and tf.gradients(y[0], x) do).