Add jax dispatch for searchsorted

Author: jessegrabowski

pytensor/link/jax/dispatch/extra_ops.py

@@ -10,6 +10,7 @@
     FillDiagonalOffset,
     RavelMultiIndex,
     Repeat,
+    SearchsortedOp,
     Unique,
     UnravelIndex,
 )
@@ -130,3 +131,13 @@ def jax_funcify_FillDiagonalOffset(op, **kwargs):
     # return filldiagonaloffset
 
     raise NotImplementedError("flatiter not implemented in JAX")
+
+
+@jax_funcify.register(SearchsortedOp)
+def jax_funcify_SearchsortedOp(op, **kwargs):
+    side = op.side
+
+    def searchsorted(x, v, side=side, sorter=None):
+        return jnp.searchsorted(x=x, v=v, side=side, sorter=sorter)
+
+    return searchsorted

tests/link/jax/test_extra_ops.py

@@ -55,6 +55,14 @@ def test_extra_ops():
         fgraph, [get_test_value(i) for i in fgraph.inputs], must_be_device_array=False
     )
 
+    values = np.arange(10)
+    query = np.array(6)
+    out = pt_extra_ops.searchsorted(values, query)
+    fgraph = FunctionGraph([], out)
+    compare_jax_and_py(
+        fgraph, [get_test_value(i) for i in fgraph.inputs], must_be_device_array=False
+    )
+
 
 @pytest.mark.xfail(reason="Jitted JAX does not support dynamic shapes")
 def test_bartlett_dynamic_shape():