Repository: google/jax

Issue:
Implement i0e and i1e bessel functions
These are useful for the von-Mises distribution on the unit circle.


Discussion:
I'm hoping you can help us with these :) I didn't see any mention of these in TF2XLA; is there an existing implementation we can draw on? Or should we treat this as a CustomCall on CPU and GPU (like we treat linalg)?

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -180,6 +180,18 @@ def digamma(x):
   r"""Elementwise digamma: :math:`\psi(x)`."""
   return digamma_p.bind(x)
 
+def bessel_i0e(x):
+  r"""Exponentially scaled modified Bessel function of order 0:
+  :math:`\mathrm{i0e}(x) = e^{-\mathrm{abs}(x)} \mathrm{i0}(x)`
+  """
+  return bessel_i0e_p.bind(x)
+
+def bessel_i1e(x):
+  r"""Exponentially scaled modified Bessel function of order 1:
+  :math:`\mathrm{i1e}(x) = e^{-\mathrm{abs}(x)} \mathrm{i1}(x)`
+  """
+  return bessel_i1e_p.bind(x)
+
 def erf(x):
   r"""Elementwise error function: :math:`\mathrm{erf}(x)`."""
   return erf_p.bind(x)
@@ -1611,6 +1623,19 @@ def _brcast_to(x, shape):
 
 digamma_p = standard_unop(_float, 'digamma')
 
+bessel_i0e_p = standard_unop(_float, 'bessel_i0e')
+ad.defjvp2(bessel_i0e_p, lambda g, y, x: g * (bessel_i1e(x) - sign(x) * y))
+
+bessel_i1e_p = standard_unop(_float, 'bessel_i1e')
+def _bessel_i1e_jvp(g, y, x):
+  eps = onp.finfo(_dtype(x)).eps
+  x_is_not_tiny = abs(x) > eps
+  safe_x = select(x_is_not_tiny, x, full_like(x, eps))
+  dy_dx = bessel_i0e(safe_x) - y * (sign(safe_x) + reciprocal(safe_x))
+  dy_dx = select(x_is_not_tiny, dy_dx, full_like(x, 0.5))
+  return g * dy_dx
+ad.defjvp2(bessel_i1e_p, _bessel_i1e_jvp)
+
 erf_p = standard_unop(_float, 'erf')
 ad.defjvp(erf_p, lambda g, x: mul(_const(x, 2. / onp.sqrt(onp.pi)),
                                   mul(g, exp(neg(square(x))))))
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -521,3 +521,11 @@ def _norm_logpdf(x):
 
 defjvp(log_ndtr,
        lambda g, ans, x: lax.mul(g, lax.exp(lax.sub(_norm_logpdf(x), ans))))
+
+@_wraps(osp_special.i0e)
+def i0e(x):
+  return lax.bessel_i0e(x)
+
+@_wraps(osp_special.i1e)
+def i1e(x):
+  return lax.bessel_i1e(x)
\ No newline at end of file


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -70,6 +70,8 @@ def op_record(name, nargs, dtypes, rng, test_grad, test_name=None):
     op_record("ndtr", 1, float_dtypes, jtu.rand_default(), True),
     # TODO(phawkins): gradient of entr yields NaNs.
     op_record("entr", 1, float_dtypes, jtu.rand_default(), False),
+    op_record("i0e", 1, float_dtypes, jtu.rand_default(), True),
+    op_record("i1e", 1, float_dtypes, jtu.rand_default(), True),
 ]
 
 CombosWithReplacement = itertools.combinations_with_replacement
@@ -102,16 +104,17 @@ def lax_fun(array_to_reduce):
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lax_fun, args_maker, check_dtypes=True)
 
-  @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": jtu.format_test_name_suffix(
-          rec.test_name, shapes, dtypes),
-       "rng": rec.rng, "shapes": shapes, "dtypes": dtypes,
-       "test_autodiff": rec.test_autodiff,
-       "scipy_op": getattr(osp_special, rec.name),
-       "lax_op": getattr(lsp_special, rec.name)}
-      for rec in JAX_SPECIAL_FUNCTION_RECORDS
-      for shapes in CombosWithReplacement(all_shapes, rec.nargs)
-      for dtypes in CombosWithReplacement(rec.dtypes, rec.nargs)))
+  @parameterized.named_parameters(itertools.chain.from_iterable(
+    jtu.cases_from_list(
+        {"testcase_name": jtu.format_test_name_suffix(
+            rec.test_name, shapes, dtypes),
+         "rng": rec.rng, "shapes": shapes, "dtypes": dtypes,
+         "test_autodiff": rec.test_autodiff,
+         "scipy_op": getattr(osp_special, rec.name),
+         "lax_op": getattr(lsp_special, rec.name)}
+        for shapes in CombosWithReplacement(all_shapes, rec.nargs)
+        for dtypes in CombosWithReplacement(rec.dtypes, rec.nargs))
+      for rec in JAX_SPECIAL_FUNCTION_RECORDS))
   def testScipySpecialFun(self, scipy_op, lax_op, rng, shapes, dtypes,
                           test_autodiff):
     args_maker = self._GetArgsMaker(rng, shapes, dtypes)