Repository: google/jax

Issue:
feat (scipy.special): Add hyp1f1 confluent hypergeometric function
I would like to implement a logpdf for the sum of n=2 independent gamma distributions for pymc.
For the special case of n=2 Mathai's Method can be used like Eqn. 3 in [https://doi.org/10.1007/s00180-019-00924-9](https://doi.org/10.1007/s00180-019-00924-9) which involves the confluent hypergeometric function $\prescript{}{1}{\mathbf{F}}{_1}(a,b,z)$. 
- scipy.special ships with the hyp1f1 function.
- The C source code can be found [here](https://github.com/scipy/scipy/blob/main/scipy/special/cephes/hyperg.c).

What would be the way to go? Translate the code into pure Python/Numpy and let jax handle it with jit?
I think Tensorflow's XLA has no variant of the function, but one for the (~~confluent~~) hypergeometric function $\prescript{}{2}{\mathbf{F}}{_1}(a,b,z)$ [cf. here](https://www.tensorflow.org/probability/api_docs/python/tfp/math/hypergeometric/hyp2f1_small_argument)

Are there any plans (maybe also considering XLA developement) to add such a functionality?



Discussion:
This would be useful for me too! Did you find an alternative in the meantime?
For me as well! +1
(It's also related to [#2991](https://github.com/google/jax/issues/2991), given that the two functions (and the way their derivatives are computed) are quite similar.)
I think that the implementation proposed by [Gil & al 2023](https://link.springer.com/article/10.1007/s11075-023-01515-y) might be suited for JAX. I'll try to translate it in LAX/JAX in the next days 

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -1719,3 +1719,183 @@ def bernoulli(n: int) -> Array:
   k = jnp.arange(2, 50, dtype=bn.dtype)  # Choose 50 because 2 ** -50 < 1E-15
   q2 = jnp.sum(k[:, None] ** -m[None, :], axis=0)
   return bn.at[4::2].set(q1 * (1 + q2))
+
+
+@custom_derivatives.custom_jvp
+@_wraps(osp_special.poch, module='scipy.special', lax_description="""\
+The JAX version only accepts positive and real inputs.""")
+def poch(z: ArrayLike, m: ArrayLike) -> Array:
+  # Factorial definition when m is close to an integer, otherwise gamma definition.
+  z, m = promote_args_inexact("poch", z, m)
+
+  return jnp.where(m == 0., jnp.array(1, dtype=z.dtype), gamma(z + m) / gamma(z))
+
+
+def _poch_z_derivative(z, m):
+  """
+  Defined in :
+  https://functions.wolfram.com/GammaBetaErf/Pochhammer/20/01/01/
+  """
+
+  return (digamma(z + m) - digamma(z)) * poch(z, m)
+
+
+def _poch_m_derivative(z, m):
+  """
+  Defined in :
+  https://functions.wolfram.com/GammaBetaErf/Pochhammer/20/01/02/
+  """
+
+  return digamma(z + m) * poch(z, m)
+
+
+poch.defjvps(
+  lambda z_dot, primal_out, z, m:  _poch_z_derivative(z, m) * z_dot,
+  lambda m_dot, primal_out, z, m: _poch_m_derivative(z, m) * m_dot,
+)
+
+
+def _hyp1f1_serie(a, b, x):
+  """
+  Compute the 1F1 hypergeometric function using the taylor expansion
+  See Eq. 3.2 and associated method (a) from PEARSON, OLVER & PORTER 2014
+  https://doi.org/10.48550/arXiv.1407.7786
+  """
+
+  def body(state):
+    serie, k, term = state
+    serie += term
+    term *= (a + k) / (b + k) * x / (k + 1)
+    k += 1
+
+    return serie, k, term
+
+  def cond(state):
+    serie, k, term = state
+
+    return (k < 250) & (lax.abs(term) / lax.abs(serie) > 1e-8)
+
+  init = 1, 1, a / b * x
+
+  return lax.while_loop(cond, body, init)[0]
+
+
+def _hyp1f1_asymptotic(a, b, x):
+  """
+  Compute the 1F1 hypergeometric function using asymptotic expansion
+  See Eq. 3.8 and simplification for real inputs from PEARSON, OLVER & PORTER 2014
+  https://doi.org/10.48550/arXiv.1407.7786
+  """
+
+  def body(state):
+    serie, k, term = state
+    serie += term
+    term *= (b - a + k) * (1 - a + k) / (k + 1) / x
+    k += 1
+
+    return serie, k, term
+
+  def cond(state):
+    serie, k, term = state
+
+    return (k < 250) & (lax.abs(term) / lax.abs(serie) > 1e-8)
+
+  init = 1, 1, (b - a) * (1 - a) / x
+  serie = lax.while_loop(cond, body, init)[0]
+
+  return gamma(b) / gamma(a) * lax.exp(x) * x ** (a - b) * serie
+
+
+@jit
+@jnp.vectorize
+def _hyp1f1_a_derivative(a, b, x):
+  """
+  Define it as a serie using :
+  https://functions.wolfram.com/HypergeometricFunctions/Hypergeometric1F1/20/01/01/
+  """
+
+  def body(state):
+    serie, k, term = state
+    serie += term * (digamma(a + k) - digamma(a))
+    term *= (a + k) / (b + k) * x / (k + 1)
+    k += 1
+
+    return serie, k, term
+
+  def cond(state):
+    serie, k, term = state
+
+    return (k < 250) & (lax.abs(term) / lax.abs(serie) > 1e-15)
+
+  init = 0, 1, a / b * x
+
+  return lax.while_loop(cond, body, init)[0]
+
+
+@jit
+@jnp.vectorize
+def _hyp1f1_b_derivative(a, b, x):
+  """
+  Define it as a serie using :
+  https://functions.wolfram.com/HypergeometricFunctions/Hypergeometric1F1/20/01/02/
+  """
+
+  def body(state):
+    serie, k, term = state
+    serie += term * (digamma(b) - digamma(b + k))
+    term *= (a + k) / (b + k) * x / (k + 1)
+    k += 1
+
+    return serie, k, term
+
+  def cond(state):
+    serie, k, term = state
+
+    return (k < 250) & (lax.abs(term) / lax.abs(serie) > 1e-15)
+
+  init = 0, 1, a / b * x
+
+  return lax.while_loop(cond, body, init)[0]
+
+
+@jit
+def _hyp1f1_x_derivative(a, b, x):
+  """
+  Define it as a serie using :
+  https://functions.wolfram.com/HypergeometricFunctions/Hypergeometric1F1/20/01/04/
+  """
+
+  return a / b * hyp1f1(a + 1, b + 1, x)
+
+
+@custom_derivatives.custom_jvp
+@jit
+@jnp.vectorize
+@_wraps(osp_special.hyp1f1, module='scipy.special', lax_description="""\
+The JAX version only accepts positive and real inputs. Values of a, b and x
+leading to high values of 1F1 might be erroneous, considering enabling double
+precision. Convention for a = b = 0 is 1, unlike in scipy's implementation.""")
+def hyp1f1(a, b, x):
+  """
+  Implementation of the 1F1 hypergeometric function for real valued inputs
+  Backed by https://doi.org/10.48550/arXiv.1407.7786
+  There is room for improvement in the implementation using recursion to
+  evaluate lower values of hyp1f1 when a or b or both are > 60-80
+  """
+  a, b, x = promote_args_inexact('hyp1f1', a, b, x)
+
+  result = lax.cond(lax.abs(x) < 100, _hyp1f1_serie, _hyp1f1_asymptotic, a, b, x)
+  index = (a == 0) * 1 + ((a == b) & (a != 0)) * 2 + ((b == 0) & (a != 0)) * 3
+
+  return lax.select_n(index,
+                      result,
+                      jnp.array(1, dtype=x.dtype),
+                      jnp.exp(x),
+                      jnp.array(jnp.inf, dtype=x.dtype))
+
+
+hyp1f1.defjvps(
+  lambda a_dot, primal_out, a, b, x: _hyp1f1_a_derivative(a, b, x) * a_dot,
+  lambda b_dot, primal_out, a, b, x: _hyp1f1_b_derivative(a, b, x) * b_dot,
+  lambda x_dot, primal_out, a, b, x: _hyp1f1_x_derivative(a, b, x) * x_dot
+)
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -54,4 +54,6 @@
   zeta as zeta,
   kl_div as kl_div,
   rel_entr as rel_entr,
+  poch as poch,
+  hyp1f1 as hyp1f1,
 )


Test patch:
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -141,7 +141,8 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
     op_record(
         "rel_entr", 2, float_dtypes, jtu.rand_positive, True,
     ),
-
+    op_record("poch", 2, float_dtypes, jtu.rand_positive, True),
+    op_record("hyp1f1", 3, float_dtypes, functools.partial(jtu.rand_uniform, low=0.5, high=30), True)
 ]