jax vmap searchsorted

PythonOT · Oct 16, 2024 · 2331d19 · 2331d19
1 parent 4862b2c
commit 2331d19
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ot/backend.py b/ot/backend.py
@@ -1590,7 +1590,7 @@ def searchsorted(self, a, v, side='left'):
  if a.ndim == 1:
  return jnp.searchsorted(a, v, side)
  else:
- return jax.vmap(jnp.searchsorted, in_axes=[0, 1, None])(a, v, side)
+ return jax.vmap(lambda b, u: jnp.searchsorted(b, u, side))(a, v)
 
  def flip(self, a, axis=None):
  return jnp.flip(a, axis)