13

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)?

vlsd
  • 945
  • 6
  • 18
  • Einsum has nothing to do with dot or tensordot, its completely utterly distinct. you could do einsum logic by hand, cast to another type if possible or implement einsum for object types. Actually implementing an `object_einsum` should not even be very hard... – seberg Mar 25 '13 at 10:17
  • @seberg, by `completely utterly distinct` you refer to the implementations, and not the mathematical notions, right? Is there any reason they are implemented differently? I can see `einsum` being implemented in terms of `tensordot` for instance, but maybe I am missing something. – Krastanov Mar 25 '13 at 11:14

3 Answers3

6

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.

FiatLux
  • 163
  • 1
  • 3
  • I wonder if this worked 7 years ago. Then again, I'm afraid of finding out it did and I essentially "wasted" a bunch of time learning how to use `nditer` – vlsd Mar 01 '21 at 20:07
  • It seems to work for some cases, but not all. For instance, summing rows of a 2D tensor still fails with this flag. – David Ketcheson Mar 28 '21 at 09:54
4

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...

seberg
  • 8,785
  • 2
  • 31
  • 30
  • Thanks, I'll give it a try! I am not quite sure I understand your point about C though... I would have to have access to the objects' mult() function in order to know how to multiply them, and I'm kind of a newb at this sort of stuff. Would you mind expanding on that? – vlsd Mar 25 '13 at 14:42
  • 1
    Yes, of course you need to use python C-API, but you can read that up. What I mean is that the nditer part itself is almost the same with `AdvancedNew` and can be almost fully copied from the existing einsum function. – seberg Mar 25 '13 at 14:57
  • @seberg, I am comparing the two answers. Is there a reason to use the deeper numpy APIs as you do instead of just rewriting `einsum` in terms of `tensordot`. Am I missing some correctness checks? Is there something besides correctness checks done better in your code? Is it performance (after all we work with `object` which is rarely performing well). – Krastanov Mar 25 '13 at 16:37
  • 2
    @Krastanov, mine is what einsum actually does, and I think nditer is pretty damn elegant. Your code does not cover everything einsum does (repeated indices), otherwise it is likely somewhat faster if you care about that. But translate this into C, and it is by far the best solution. – seberg Mar 25 '13 at 16:59
2

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

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.

Krastanov
  • 6,479
  • 3
  • 29
  • 42