jax-ml/jax · error · NotImplementedError

Negative orders for normalization is not implemented yet.

Error message

Negative orders for normalization is not implemented yet.

What it means

lpmn (associated Legendre functions) with is_normalized=True computes derivatives via _gen_derivatives, but the derivative recurrence requires negative orders, and normalized negative-order Legendre functions are not implemented in JAX. Hence the NotImplementedError.

Source

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

  Returns:
    The 3D array representing the derivatives of associated Legendre functions
    of the first kind.
  """

  num_m, num_l, num_x = p.shape

  # p_{l-1}^m.
  p_m_lm1 = jnp.pad(p, ((0, 0), (1, 0), (0, 0)))[:, :num_l, :]

  # p_{l-1}^{m+2}.
  p_mp2_lm1 = jnp.pad(p_m_lm1, ((0, 2), (0, 0), (0, 0)))[2:num_m + 2, :, :]

  # p_{l-1}^{m-2}.
  p_mm2_lm1 = jnp.pad(p_m_lm1, ((2, 0), (0, 0), (0, 0)))[:num_m, :, :]

  # Derivative computation requires negative orders.
  if is_normalized:
    raise NotImplementedError(
        'Negative orders for normalization is not implemented yet.')
  else:
    if num_l > 1:
      l_vec = jnp.arange(1, num_l - 1, dtype=x.dtype)
      p_p1 = p[1, 1:num_l - 1, :]
      coeff = -1.0 / ((l_vec + 1) * l_vec)
      update_p_p1 = jnp_einsum.einsum('i,ij->ij', coeff, p_p1)
      p_mm2_lm1 = p_mm2_lm1.at[1, 2:num_l, :].set(update_p_p1)

    if num_l > 2:
      l_vec = jnp.arange(2, num_l - 1, dtype=x.dtype)
      p_p2 = p[2, 2:num_l - 1, :]
      coeff = 1.0 / ((l_vec + 2) * (l_vec + 1) * l_vec * (l_vec - 1))
      update_p_p2 = jnp_einsum.einsum('i,ij->ij', coeff, p_p2)
      p_mm2_lm1 = p_mm2_lm1.at[0, 3:num_l, :].set(update_p_p2)

  m_mat, l_mat = jnp.meshgrid(
    jnp.arange(num_m, dtype=x.dtype),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use is_normalized=False: lpmn(m, n, z, is_normalized=False, diff=True) works because the unnormalized branch implements negative orders
2. Normalize manually after computing unnormalized values/derivatives using known normalization factors
3. Compute gradients numerically (finite differences) for the normalized variant as a stopgap

Example fix

// before
p, dp = jax.scipy.special.lpmn(m, n, z, is_normalized=True)  # diff default
// after
p, dp = jax.scipy.special.lpmn(m, n, z, is_normalized=False)
# apply normalization factors yourself if needed

if is_normalized and diff:
    use_is_normalized = False  # fall back to unnormalized + manual normalization

try:
    p, dp = lpmn(m, n, z, is_normalized=True)
except NotImplementedError:
    p, dp = lpmn(m, n, z, is_normalized=False)
    p, dp = normalize(p), normalize(dp)

When it happens

Trigger: Calling jax.scipy.special.lpmn(m, n, z, is_normalized=True) with diff=True (the default), which routes into _gen_derivatives and hits the normalized branch.

Common situations: Porting scipy.special.lpmn-based geophysics or quantum mechanics code (spherical harmonics gradients) that uses normalized Legendre polynomials and their derivatives; gradient-based optimization over spherical-harmonic coefficients.

Related errors

• z must be a 1D array.
• Riemann zeta function not implemented; pass q != None to com
• z.dtype={} is not supported, see docstring for supported typ
• Computations for m!=n are not yet supported.
• The 'diff_n' argument to jax.scipy.special.sph_harm_y is not

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