我如何让 numpy.einsum 与 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.

Answer 3

Here is a much simpler implementation that separates the einsum in multiple tensordot s.

def einsum(string, *args):
    index_groups = map(list, string.split(','))
    assert len(index_groups) == len(args)
    tensor_indices_tuples = zip(index_groups, args)
    return reduce(einsum_for_two, tensor_indices_tuples)[1]

def einsum_for_two(tensor_indices1, tensor_indices2):
    string1, tensor1 = tensor_indices1
    string2, tensor2 = tensor_indices2
    sum_over_indices = set(string1).intersection(set(string2))
    new_string = string1 + string2
    axes = ([], [])
    for i in sum_over_indices:
        new_string.remove(i)
        new_string.remove(i)
        axes[0].append(string1.index(i))
        axes[1].append(string2.index(i))
    return new_string, np.tensordot(tensor1, tensor2, axes)

First it separates the einsum arguments in tuples of (indices, tensor). Then it reduces of the list as follows:

• Takes the first two tuples, and evaluates a simple einsum_for_two on them. It also prints out the new indices signature.
• The value of einsum_for_two is used with the next tuple in the list as the new arguments for einsum_for_two .
• Continues until there is only tuple left. The indices signature is discarded and only the tensor is returned.

It is probably slow (but anyway you are using object dtype ). It does not do many correctness checks on the input.

As @seberg noted, my code does not work for traces of tensors.