implement hyp2f1

mattbahr · mattbahr · commit 28cbf8ddd9a2 · 2025-04-22T15:57:23.000Z
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
@@ -2637,6 +2637,341 @@ def hyp1f1(a: ArrayLike, b: ArrayLike, x: ArrayLike) -> Array:
 )
 
 
+def _binom(n, k):
+  a = lax.lgamma(n + 1.0)
+  b = lax.lgamma(n - k + 1.0)
+  c = lax.lgamma(k + 1.0)
+
+  return lax.exp(a - b - c)
+
+
+def _poch(q, n):
+  """
+  `jax.scipy.special.poch` does not allow for non-positive integer q.
+  """
+  def body(i, state):
+    q, prod = state
+
+    prod *= q + i
+
+    return q, prod
+
+  return lax.cond(
+    n == 0,
+    lambda: jnp.array(1, dtype=q.dtype),
+    lambda: lax.fori_loop(jnp.array(1, dtype=n.dtype), n, body, (q, q))[1]
+  )
+
+
+def _hyp2f1_terminal(a, b, c, x):
+  """
+  The Taylor series representation of the 2F1 hypergeometric function
+  terminates when either a or b is a non-positive integer. See Eq. 4.1 and
+  Taylor Series Method (a) from PEARSON, OLVER & PORTER 2014
+  https://doi.org/10.48550/arXiv.1407.7786
+  """
+  # Ensure that between a and b, the negative integer parameter with the greater
+  # absolute value - that still has a magnitude less than the absolute value of
+  # c if c is non-positive - is used for the upper limit in the loop.
+  temp = a
+  a = jnp.where(
+    jnp.logical_and(
+      b < a,
+      jnp.logical_and(
+        b % 1 == 0,
+        jnp.logical_not(
+          jnp.logical_and(
+            c % 1 == 0,
+            jnp.logical_and(
+              c <= 0,
+              c > b
+            )
+          )
+        )
+      )
+    ), b, a
+  )
+  b = jnp.where(
+    jnp.logical_and(
+      b < temp,
+      jnp.logical_and(
+        b % 1 == 0,
+        jnp.logical_not(
+          jnp.logical_and(
+            c % 1 == 0,
+            jnp.logical_and(
+              c <= 0,
+              c > b
+            )
+          )
+        )
+      )
+    ), temp, b
+  )
+
+  def body(i, sum):
+    sum += (-1) ** i * _binom(jnp.abs(a), i) / _poch(c, i) * _poch(b, i) * x ** i
+
+    return sum
+
+  return lax.fori_loop(jnp.array(0, dtype=a.dtype),
+                       jnp.abs(a) + 1,
+                       body,
+                       jnp.array(0, dtype=x.dtype))
+
+
+def _hyp2f1_serie(a, b, c, x):
+  """
+  Compute the 2F1 hypergeometric function using the Taylor expansion.
+  See Eq. 4.1 from PEARSON, OLVER & PORTER 2014
+  https://doi.org/10.48550/arXiv.1407.7786
+  """
+  precision = jnp.finfo(jnp.float32).eps
+
+  s = 1 - x
+
+  neg_int_a = jnp.logical_and(a <= 0, a % 1 == 0)
+  neg_int_b = jnp.logical_and(b <= 0, b % 1 == 0)
+  neg_int_c = jnp.logical_and(c <= 0, c % 1 == 0)
+
+  def body(state):
+    serie, k, term = state
+    serie += term
+    term = _poch(a, k) / _poch(c, k) * _poch(b, k) / factorial(k) * x ** k
+    k += 1
+
+    return serie, k, term
+
+  def cond(state):
+    serie, k, term = state
+
+    return (k < 250) & (lax.abs(term) / lax.abs(serie) > precision)
+
+  init = (jnp.array(0, dtype=x.dtype),
+          jnp.array(1, dtype=x.dtype),
+          jnp.array(1, dtype=x.dtype))
+
+  return lax.while_loop(cond, body, init)[0]
+
+
+def _hyp2f1_terminal_or_serie(a, b, c, x):
+  """
+  Check for recurrence relations along with whether or not the series
+  terminates. True recursion is not possible; however, the recurrence
+  relation may still be approximated.
+  See 4.6.1. Recurrence Relations from PEARSON, OLVER & PORTER 2014
+  https://doi.org/10.48550/arXiv.1407.7786
+  """
+  neg_int_a = jnp.logical_and(a <= 0, a % 1 == 0)
+  neg_int_b = jnp.logical_and(b <= 0, b % 1 == 0)
+  neg_int_c = jnp.logical_and(c <= 0, c % 1 == 0)
+  neg_int_a_or_b = jnp.logical_or(neg_int_a, neg_int_b)
+  not_neg_int_a_or_b = jnp.logical_not(neg_int_a_or_b)
+
+  s = 1 - x
+  d = c - a - b
+
+  index = jnp.where(
+    jnp.logical_and(
+      neg_int_c,
+      jnp.logical_and(
+        jnp.logical_not(jnp.logical_and(neg_int_a, a > c)),
+        jnp.logical_not(jnp.logical_and(neg_int_b, b > c))
+      )
+    ), 0,
+    jnp.where(jnp.logical_and(x < -0.5, not_neg_int_a_or_b),
+      jnp.where(b > a, 1, 2),
+      jnp.where(jnp.logical_and(x > 0.9, not_neg_int_a_or_b),
+        jnp.where(d % 1 != 0, 3, 4),
+        jnp.where(jnp.logical_and(jnp.logical_not(neg_int_c), neg_int_a_or_b), 5, 3))))
+
+  return lax.select_n(index,
+                      jnp.array(jnp.inf, dtype=x.dtype),
+                      s ** (-a) * _hyp2f1_serie(a, c - b, c, -x / s),
+                      s ** (-b) * _hyp2f1_serie(c - a, b, c, -x / s),
+                      _hyp2f1_serie(a, b, c, x),
+                      _hyp2f1_digamma_transform(a, b, c, x),
+                      _hyp2f1_terminal(a, b, c, x))
+
+
+def _hyp2f1_gamma_transform(a, b, c, x):
+  """
+  Gamma transformations of the 2F1 hypergeometric function.
+  """
+
+  def transform_1():
+    """
+    See Eq. 4.10 and Analytic Continuation Formulas from PEARSON, OLVER & PORTER 2014
+    https://doi.org/10.48550/arXiv.1407.7786
+    """
+    p = _hyp2f1_serie(a, 1 - c + a, 1 - b + a, 1 / x)
+    q = _hyp2f1_serie(b, 1 - c + b, 1 - a + b, 1 / x)
+    p *= (-x) ** (-a)
+    q *= (-x) ** (-b)
+    t1 = gamma(c)
+    s = t1 * gamma(b - a) / (gamma(b) * gamma(c - a))
+    y = t1 * gamma(a - b) / (gamma(a) * gamma(c - b))
+
+    return s * p + y * q
+
+  def transform_2():
+    """
+    See 4.1 Properties of F from PEARSON, OLVER & PORTER 2014
+    https://doi.org/10.48550/arXiv.1407.7786
+    """
+    return gamma(c) * gamma(c - a - b) / (gamma(c - a) * gamma(c - b))
+
+  return jnp.where(
+    x < -2,
+    transform_1(),
+    transform_2()
+  )
+
+
+def _hyp2f1_digamma_transform(a, b, c, x):
+  """
+  Digamma transformation of the 2F1 hypergeometric function.
+  See AMS55 #15.3.10, #15.3.11, #15.3.12
+  """
+  precision = jnp.finfo(jnp.float32).eps
+
+  d = c - a - b
+  s = 1 - x
+  id = jnp.round(d)
+
+  e = jnp.where(id >= 0, d, -d)
+  d1 = jnp.where(id >= 0, d, jnp.array(0, dtype=d.dtype))
+  d2 = jnp.where(id >= 0, jnp.array(0, dtype=d.dtype), d)
+  aid = jnp.where(id >= 0, id, -id).astype('int32')
+
+  ax = jnp.log(s)
+
+  y = digamma(1.0) + digamma(1.0 + e) - digamma(a + d1) - digamma(b + d1) - ax
+  y /= gamma(e + 1.0)
+
+  p = (a + d1) * (b + d1) * s / gamma(e + 2.0)
+
+  def cond(state):
+    _, _, _, _, _, _, q, _, _, t, y = state
+
+    return jnp.logical_and(
+      t < 250,
+      jnp.logical_or(y == 0, jnp.abs(q / y) > precision)
+    )
+
+  def body(state):
+    a, ax, b, d1, e, p, q, r, s, t, y = state
+
+    r = digamma(1.0 + t) + digamma(1.0 + t + e) - digamma(a + t + d1) \
+        - digamma(b + t + d1) - ax
+    q = p * r
+    y += q
+    p *= s * (a + t + d1) / (t + 1.0)
+    p *= (b + t + d1) / (t + 1.0 + e)
+    t += 1.0
+
+    return a, ax, b, d1, e, p, q, r, s, t, y
+
+  init = (a, ax, b, d1, e, p, y, jnp.array(0, dtype=x.dtype), s,
+          jnp.array(1, dtype=x.dtype), y)
+  _, _, _, _, _, _, q, r, _, _, y = lax.while_loop(cond, body, init)
+
+  def compute_sum(y):
+    y1 = jnp.array(1, dtype=x.dtype)
+    t = jnp.array(0, dtype=x.dtype)
+    p = jnp.array(1, dtype=x.dtype)
+
+    def for_body(i, state):
+      a, b, d2, e, p, s, t, y1 = state
+
+      r = 1.0 - e + t
+      p *= s * (a + t + d2) * (b + t + d2) / r
+      t += 1.0
+      p /= t
+      y1 += p
+
+      return a, b, d2, e, p, s, t, y1
+
+    init_val = a, b, d2, e, p, s, t, y1
+    y1 = lax.fori_loop(1, aid, for_body, init_val)[-1]
+
+    p = gamma(c)
+    y1 *= gamma(e) * p / (gamma(a + d1) * gamma(b + d1))
+    y *= p / (gamma(a + d2) * gamma(b + d2))
+
+    y = jnp.where((aid & 1) != 0, -y, y)
+    q = s ** id
+
+    return jnp.where(id > 0, y * q + y1, y + y1 * q)
+
+  return jnp.where(
+    id == 0,
+    y * gamma(c) / (gamma(a) * gamma(b)),
+    compute_sum(y)
+  )
+
+
+@jit
+@jnp.vectorize
+def hyp2f1(a: ArrayLike, b: ArrayLike, c: ArrayLike, x: ArrayLike) -> Array:
+  r"""The 2F1 hypergeometric function.
+
+  JAX implementation of :obj:`scipy.special.hyp2f1`.
+
+  .. math::
+
+     \mathrm{hyp2f1}(a, b, c, x) = {}_2F_1(a; b; c; x) = \sum_{k=0}^\infty \frac{(a)_k(b)_k}{(c)_k}\frac{x^k}{k!}
+
+  where :math:`(\cdot)_k` is the Pochammer symbol.
+
+  The JAX version only accepts positive and real inputs. Values of
+  ``a``, ``b``, ``c``, and ``x`` leading to high values of 2F1 may
+  lead to erroneous results; consider enabling double precision in this case.
+
+  Args:
+    a: arraylike, real-valued
+    b: arraylike, real-valued
+    c: arraylike, real-valued
+    x: arraylike, real-valued
+
+  Returns:
+    array of 2F1 values.
+  """
+  # This is backed by https://doi.org/10.48550/arXiv.1407.7786
+  a, b, c, x = promote_args_inexact('hyp2f1', a, b, c, x)
+
+  d = c - a - b
+  s = 1 - x
+
+  neg_int_ca = jnp.logical_and(c - a <= 0, (c - a) % 1 == 0)
+  neg_int_cb = jnp.logical_and(c - b <= 0, (c - b) % 1 == 0)
+  neg_int_ca_or_cb = jnp.logical_or(neg_int_ca, neg_int_cb)
+
+  index = jnp.where(jnp.logical_or(x == 0, jnp.logical_and(jnp.logical_or(a == 0, b == 0), c != 0)), 0,
+            jnp.where(c == 0, 2,
+              jnp.where(jnp.logical_and(d <= -1, jnp.logical_not(jnp.logical_and(d % 1 != 0, s < 0))), 1,
+                jnp.where(jnp.logical_and(d <= 0, x == 1), 2,
+                  jnp.where(jnp.logical_and(x < 1, b == c), 3,
+                    jnp.where(jnp.logical_and(x < 1, a == c), 4,
+                      jnp.where(x > 1, 2,
+                        jnp.where(x == 1,
+                          jnp.where(neg_int_ca_or_cb,
+                            jnp.where(d >= 0, 5, 2),
+                            jnp.where(d <= 0, 2, 6)),
+                          jnp.where(d < 0, 7,
+                            jnp.where(neg_int_ca_or_cb, 5, 7))))))))))
+
+  return lax.select_n(index,
+                      jnp.array(1, dtype=x.dtype),
+                      s ** d * _hyp2f1_terminal_or_serie(c - a, c - b, c, x),
+                      jnp.array(jnp.inf, dtype=x.dtype),
+                      s ** (-a),
+                      s ** (-b),
+                      s ** d * _hyp2f1_serie(c - a, c - b, c, x),
+                      _hyp2f1_gamma_transform(a, b, c, x),
+                      _hyp2f1_terminal_or_serie(a, b, c, x))
+
+
 def softmax(x: ArrayLike,
             /,
             *,
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
@@ -37,6 +37,7 @@
   gammaln as gammaln,
   gammasgn as gammasgn,
   hyp1f1 as hyp1f1,
+  hyp2f1 as hyp2f1,
   i0 as i0,
   i0e as i0e,
   i1 as i1,
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
@@ -157,6 +157,10 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
         "hyp1f1", 3, float_dtypes,
         functools.partial(jtu.rand_uniform, low=0.5, high=30), True
     ),
+    op_record(
+        "hyp2f1", 4, float_dtypes,
+        functools.partial(jtu.rand_uniform, low=0.5, high=30), False
+    ),
     op_record("log_softmax", 1, float_dtypes, jtu.rand_default, True),
     op_record("softmax", 1, float_dtypes, jtu.rand_default, True),
 ]
