Einsum wrapper in PyTensor to mimic behaviour of np.einsum

[DavidRavnsborg]

Hello,

I would very much like to see einsum added to pytensor functions. I am currently stuck trying to fit a matrix of coefficients that are multiplied by a 3D tensor, and it would be ideal to be able to go:

contribution = pt.einsum("ijk,kl->ij", tensor_product, coef_matrix)

I currently do this with numpy, but I can't fit coefficients with numpy:

contribution = np.einsum("ijk,kl->ij", tensor_product, coef_matrix)

What would it take to add this to the library? I am on the fence between re-working my code and just creating and committing a PR for this.

[ricardoV94]

Thanks for writing up. As replied in the other thread pytensor already has einsum in more recent versions, if you bump pymc to the latest version you'll see it under pytensor.tensor.einsum.