dgasmith / opt_einsum

⚡️Optimizing einsum functions in NumPy, Tensorflow, Dask, and more with contraction order optimization.
https://dgasmith.github.io/opt_einsum/
MIT License
861 stars 68 forks source link

Add support for sympy backend (symbolic computation) #242

Open SimoneGasperini opened 2 days ago

SimoneGasperini commented 2 days ago

I think it would be very useful and powerful to include also the support for the computation with symbolic $n$-dimensional arrays based on sympy (see sympy.tensor.array module). To give an explicit example, consider the following code:

import sympy
import numpy
from opt_einsum import contract

arr1 =  sympy.Array([[[1,0],[0,0]],
                     [[0,0],[0,0]]])
arr2 = sympy.Array([[0,1],[1,0]])

output = contract('ijk, kl -> ijl', arr1, arr2)

isinstance(output, sympy.Array)  # False
isinstance(output, numpy.ndarray)  # True

This is working and it returns the correct result but the output array is a numpy.ndarray with dtype=object which is unnatural and quite inefficient for symbolic computation. It would be nice to have an implementation of the contract function able to operate on sympy.Array objects directly.

Please let me know if this is something planned for the opt_einsum project in the near-term future. Otherwise, it would be great if you could drop here any reference to other possible implementation attempts of that kind of symbolic version of the optimized tensors contraction.

SimoneGasperini commented 2 days ago

Related question on Quantum Computing Stack Exchange: Apply parametric gate to quantum state in symbolic tensor representation

jcmgray commented 2 days ago

I think it should be quite possible, one just needs to write a tensordot (and possible einsum depending on the types of contract) wrapper that calls the required sympy funtions like https://docs.sympy.org/latest/modules/tensor/array.html#sympy.tensor.array.tensorcontraction. I don't have any plans to implemented this myself however.

SimoneGasperini commented 2 days ago

Ok cool, thank you @jcmgray. I can try to draft an implementation myself as soon as I find some time. Let me know if you have any specific suggestion in order to tackle the problem in the right direction from the very beginning.