import jax.numpy as jnp 
from jax import random, grad
import ot as pot
key = random.PRNGKey(0)
B = 10
key, subkey = random.split(key)
x = random.normal(subkey, (B, 1))
key, subkey = random.split(key)
y = random.normal(subkey, (B, 1))
def loss_fn(x, y):
 costs = jnp.linalg.norm(x[:, None] - y[None, :], axis=-1)**2
 pi = pot.emd(
 jnp.ones(B) / B, 
 jnp.ones(B) / B, 
 costs)
 return jnp.sum(pi * costs)
g = grad(loss_fn)(x, y)
print(g)

Traceback (most recent call last):
 File "/home/telfaralex/Documents/phdv2/code/sinterp/sinterp/tests/test_couplings.py", line 58, in test_grad
 g = grad(loss_fn)(x, y)
 ^^^^^^^^^^^^^^^^^^^
 File "/home/telfaralex/Documents/phdv2/code/sinterp/sinterp/tests/test_couplings.py", line 51, in loss_fn
 pi = ot_fn(
 ^^^^^^
 File "/home/telfaralex/miniconda3/lib/python3.11/site-packages/ot/lp/__init__.py", line 318, in emd
 M, a, b = nx.to_numpy(M, a, b)
 ^^^^^^^^^^^^^^^^^^^^
 File "/home/telfaralex/miniconda3/lib/python3.11/site-packages/ot/backend.py", line 260, in to_numpy
 return [self._to_numpy(array) for array in arrays]
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "/home/telfaralex/miniconda3/lib/python3.11/site-packages/ot/backend.py", line 260, in <listcomp>
 return [self._to_numpy(array) for array in arrays]
 ^^^^^^^^^^^^^^^^^^^^^
 File "/home/telfaralex/miniconda3/lib/python3.11/site-packages/ot/backend.py", line 1439, in _to_numpy
 return np.array(a)
 ^^^^^^^^^^^
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[10,10]
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError