ENH: add einsum

[ev-br]

A couple of notes:

• the generic normalization machinery does not handle signatures with positional args before varargs. This is einsum and gradient, so keep the generic machinery simple and do things manually here.
• einsum signature is a bit weird even in numpy, with **kwargs, so follow that. Not sure why it is what it is, one possible reason could be that keyword-only args after varargs is a syntax error (?) E.g. def func(a, *args, *, out=None): pass.
• IIUC, pytorch does not allow the same level of control over the optimize = {False, True, ‘greedy’, ‘optimal’} argument, so silenty ignore this argument

Also while at it, add a dedicated annotation to validate the casting=... argument. This is strictly speaking extraneous to this PR, so can take it out if desired.

[lezcano]

I can just imagine the pain that was implementing this one. I left a few comments, but overall looks good.

[lezcano]

Just define it at the end of _funcs.py really and leave a note as to why.

[ev-br]

If removing imports is the goal, that won't help really. These normalizers are otherwise only used in _normalizations.py, so these need to be imported either way.

[lezcano]

What I'd like to remove is local imports. They are a massive red flag

[ev-br]

That would create an import cycle since _funcs.py imports _funcs_impl. Also from _ndarray import ndarray must be local (as it is in _normalizations).

We can of course, import these at the _funcs_impl module level and special-case them to not get dumped into the global namespace, if that's really what you think is needed?

[lezcano]

We would just need to import things from _normalizations (which we already do) and well, we would still have the ndarray import dangling there, but there's not much that can be done about that really, as we have that one all throughout the codebase.

[ev-br]

OK, done in the last commit

[lezcano]

These two normalizers don't do all that much, so let's just do the error checking directly here for conciseness.

[lezcano]

Also assert that optimize is False otherwise raise a NotImplementedError similar to order and so on.

[ev-br]

Oh. Now that I re-read https://pytorch.org/docs/stable/generated/torch.einsum.html, there is torch.backends.opt_einsum = "greedy" so we can support this :-). Question: is there a way of controlling torch.backends in a context manager, other than a try... finally block?

[lezcano]

I don't think we have a context manager for that, no. try...finally seems reasonable to me.

[ev-br]

Thanks, done in 211a461

[lezcano]

Why the [:-1]? Isn't sublistout optional? Also, don't we want to preprocess that one as well, perhaps asserting that it's an ndarray?

[ev-br]

Exactly!

• If sublistout is not given, the length of operands is even, and we pick odd-numbered elements, which are arrays.
• If sublistout is given, the length of operands is odd, we peel off the last one, and pick odd-numbered elements, which are arrays. Without [:-1], we would have picked sublistout, too --- and it's a sublist, not an array.

And, no, it's not an array really. Can contain e.g. an Ellipsis:

            assert_equal(np.einsum("...i->...", a, optimize=do_opt),
                         np.sum(a, axis=-1).astype(dtype))
            assert_equal(np.einsum(a, [Ellipsis, 0], [Ellipsis], optimize=do_opt),
                         np.sum(a, axis=-1).astype(dtype))

[lezcano]

Right. Could you please leave a comment or a link to this explanation?

[lezcano]

result_dtype is not being used.

[ev-br]

thanks, removed.

NB it is not used because the current implementation follows pytorch in that e.g.

In [48]: a = np.arange(8, dtype=np.int8)

In [49]: torch.einsum(torch.as_tensor(a), [0], []).dtype
Out[49]: torch.int64

unlike numpy where the last line is int8 (sigh)

[lezcano]

Yeah, this is because, as sum and reductions that accumulate, integeras are upcasted to int64. This function simply uses sum and bmm internally so it has the same semantics as those two.

[lezcano]

Great! What a pain of a function...

[ev-br]

Thanks Mario for the review!