jax-ml/jax · error · NotImplementedError

The 'diff_n' argument to jax.scipy.special.sph_harm_y is not

Error message

The 'diff_n' argument to jax.scipy.special.sph_harm_y is not supported.

What it means

jax.scipy.special.sph_harm_y added a diff_n keyword (for derivative output) in SciPy 1.15, but JAX's implementation only computes the harmonics themselves. Passing any non-None value for diff_n raises NotImplementedError immediately.

Source

Thrown at jax/_src/scipy/special.py:2347

    n: The degree of the harmonic; must have `n >= 0`. The standard notation for
      degree in descriptions of spherical harmonics is `l (lower case L)`. We
      use `n` here to be consistent with `scipy.special.sph_harm_y`. Return
      values for `n < 0` are undefined.
    m: The order of the harmonic; must have `|m| <= n`. Return values for
      `|m| > n` are undefined.
    theta: The polar (colatitudinal) coordinate; must be in [0, pi].
    phi: The azimuthal (longitudinal) coordinate; must be in [0, 2*pi].
    diff_n: Unsupported by JAX.
    n_max: The maximum degree `max(n)`. If the supplied `n_max` is not the true
      maximum value of `n`, the results are clipped to `n_max`. For example,
      `sph_harm_y(n=jnp.array([10]), m=jnp.array([2]), theta=theta, phi=phi, n_max=6)`
      actually returns
      `sph_harm_y(n=jnp.array([6]), m=jnp.array([2]), theta=theta, phi=phi, n_max=6)`.
  Returns:
    A 1D array containing the spherical harmonics at (n, m, theta, phi).
  """
  if diff_n is not None:
    raise NotImplementedError(
        "The 'diff_n' argument to jax.scipy.special.sph_harm_y is not supported.")

  if jnp.isscalar(theta):
    theta = jnp.array([theta])

  if n_max is None:
    n_max = np.max(n)
  n_max = core.concrete_or_error(
      int, n_max, 'The `n_max` argument of `jax.scipy.special.sph_harm_y` must '
      'be statically specified to use `sph_harm_y` within JAX transformations.')

  return _sph_harm(n, m, theta, phi, n_max)


# exponential integrals
# these algorithms are ported over from the files ei.c and expn.c in the Cephes mathematical library.
# https://fossies.org/dox/cephes-math-28/ei_8c_source.html
# https://fossies.org/dox/cephes-math-28/expn_8c_source.html

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop the diff_n argument and take gradients with jax.grad/jacfwd instead.
Compute derivative arrays with scipy.special.sph_harm_y (which supports diff_n) outside JAX.
Pin/branch your code on library: use scipy directly when diff_n is required.

Example fix

# before
y = jax.scipy.special.sph_harm_y(n, m, theta, phi, n_max=6, diff_n=1)

# after
y = jax.scipy.special.sph_harm_y(n, m, theta, phi, n_max=6)
dy = jax.jacfwd(lambda t: jax.scipy.special.sph_harm_y(n, m, t, phi, n_max=6))(theta)

Defensive patterns

Strategy: fallback

Validate before calling

# no pre-call validation possible besides not passing diff_n
kwargs = {}  # never include diff_n when calling jax.scipy.special.sph_harm_y

Try / catch

try:
    y = jax.scipy.special.sph_harm_y(n, m, theta, phi, n_max=n_max)
except NotImplementedError:
    y = jax.pure_callback(
        lambda a: scipy.special.sph_harm_y(*a, diff_n=1), ...,
        (n, m, theta, phi))

Prevention

Feature-detect: only pass diff_n when calling SciPy, never jax.
Use autodiff (jacfwd/grad) for derivatives instead of diff_n.
Pin SciPy usage in a thin adapter layer so the JAX path stays clean.

When it happens

Trigger: Calling sph_harm_y(n, m, theta, phi, n_max=..., diff_n=1) or any non-None diff_n, typically code copied from new SciPy docs.

Common situations: Upgrading SciPy-facing code to SciPy >= 1.15 where diff_n is part of the signature, then running the same call under jax.scipy.special; using diff_n for gradient computations of spherical harmonics.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/17e2ed4e59219ad5. Report an issue: GitHub.