This PR enables dispatching all pairwise calls in einsum when optimize= is turned on to batched matrix multiply (or non-batched matmul and multiply as special cases). For a short explanation of the implementation see here - https://github.com/jcmgray/einsum_bmm.

Currently falling back to c_einsum for these, when there exist batch indices, induces potentially many orders of magnitude slow downs (#22604, e.g. 7000x), particularly in the case of large or high dimensional tensors.

After fairly extensive benchmarking (below), I believe that this should be a uniform improvement across essentially all the cases when one would use optimize=True. This is partly enabled by the modern performance of matmul which in the past was not always faster than einsum. However, I know einsum has an enormous number of uses and I'd gladly accept more suggestions for benchmarks, and would understand people's hesitation!

(@dgasmith, @seberg)

Other notes:

• This replaces all the can_blas logic and tensordot calls, and simplifies everything into the bmm_einsum, which encapsulates how tensordot is implemented
• When there are no contracted indices, numpy.multiply is used for slightly better performance
• The current implementation can simply pass out and other kwargs on to mutiply and matmul, unlike tensordot, meaning einsum('ij,jk->ik', x, y, optimize=True, out=z) should write directly to z
• The fusing allows einsum(eq, ..., optimize=True) to be used when there are more than 32 indices involved (Increase maximum number of array dimensions? #5744).
• The low overhead of the actuall bmm_einsum implementation comes from caching what operations to use based on (eq, x.shape, y.shape), using @functools.lru_cache(2**12).
• the performance broadly speaking is brought in line with torch.einsum

TODO:

• update docs?
• decide cache size?
• should I lint the existing einsumfunc.py code?

Benchmarks

In the following the kernels are the following:

• no optimize: the base np.einsum(eq, x, y) i.e. c_einsum
• optimize + dot: the current np.einsum(eq, x, y, optimize=True) which calls tensordot where possible, and also induces some overhead from the potential path optimization (even though there are only two terms in the following, the equation must be checked and parsed etc.)
• optimize + bmm: the proposed np.einsum(eq, x, y, optimize=True) which calls batched matmul, and still induces some overhead from the potential path optimization
• no optimize + bmm: this calls the pairwise bmm einsum directly, and is provided simply to show what overhead comes from the bmm impl, and what from the potential path optimization

The main two to compare are optimize + dot to optimize + bmm. Overall I'd summarize the results as the following:

1. when a contraction is memory bound, all methods are pretty much equivalent
2. when a contraction involves no batch dimensions, performance of dot is retained
3. when a contraction involves batch dimensions, performance can be orders of magnitude faster

All dimensions are size n unless otherwise noted.

For batch_matmul_small the sizes are (n, 2, 2), (n, 2, 2).

For batch_matmul_large the sizes are (10, n, n), (10, n, n).

In random_extreme the shapes are taken as the following:

[
    (n, 4, 3, 4, 3, 4, 2, 4, 2, n, 2, n, 3, 2, 4),
    (2, 4, n, n, 4, 4, 4, 3, 4, 4, 4, 3, 4)
]

In many_small_dims_dot the shapes are ([2] * n, [2] * n) and half in an interleaved pattern are contracted.

In many_small_dims_batched_dot the shapes are ([2] * n, [2] * n) and a third in an interleaved pattern are contracted and a third in an interleaved pattern are batch indices.