How do I get numpy.einsum to play well with sympy?

Question

Ok, so I have several, multi-dimensional numpy arrays of sympy objects (expressions). For example:

A = array([[1.0*cos(z0)**2 + 1.0, 1.0*cos(z0)],
          [1.0*cos(z0), 1.00000000000000]], dtype=object)

and so on.

What I would like to do is multiply several of these arrays using einsum, since I already have the syntax for that from a numerical calculation I was doing earlier. The problem is, when I try to do something like

einsum('ik,jkim,j', A, B, C)

I get a type error:

TypeError: invalid data type for einsum

Sure, so a quick search on Google shows me einsum probably can't do this, but no reason as to why. In particular, calling the numpy.dot() and numpy.tensordot() functions on those arrays works like a charm. I could use tensordot to do what I need, but my brain hurts when I think about having to replace fifty or so Einsten summations like the one above (where the order of the indeces is very important) with nested tensordot calls. Even more nightmarish is the though of having to debug that code and hunt for that one misplaced index swap.

Long story short, does anyone know why tensordot works with objects but einsum will not? Any suggestions towards a workaround? If not, any suggestions as to how I would go about writing my own wrapper to nested tensordot calls that is somewhat similar to the einsum notation (numbers instead of letters are fine)?

Answer 1

Einsum basically supersedes tensordot (not dot, because dot is normally using optimized linear algebra packages), code wise it is completely different.

Here is an object einsum, its untested (for more complex things), but I think it should work... Doing the same thing in C is probably even simpler since you can steal everything but the loop itself from the real einsum function. So if you feel like it, implement it and make more people happy...

https://gist.github.com/seberg/5236560

I will not guarantee anything, especially not for weirder corner cases. Of course you can translate einsum notation to tensordot notation too I am sure, and that is probably a bit faster since the loops would end up being mostly in C...

Answer 2

Interestingly enough, adding optimize="optimal" worked for me

einsum('ik,jkim,j', A, B, C) yields error, but

einsum('ik,jkim,j', A, B, C, optimize="optimal") works perfectly well with sympy.