jax-ml/jax · error · NotImplementedError

Computations for m!=n are not yet supported.

Error message

Computations for m!=n are not yet supported.

What it means

jax.scipy.special.lpmn only implements the associated Legendre functions for the case m == n. The implementation computes all values up to degree l_max = n and assumes the m and n arguments are equal; any other combination is unimplemented in JAX (unlike SciPy). It raises NotImplementedError rather than silently returning wrong results.

Source

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

  Raises:
    TypeError if elements of array `z` are not in (float32, float64).
    ValueError if array `z` is not 1D.
    NotImplementedError if `m!=n`.
  """
  dtype = lax.dtype(z)
  if dtype not in (np.float32, np.float64):
    raise TypeError(
        'z.dtype={} is not supported, see docstring for supported types.'
        .format(dtype))

  if z.ndim != 1:
    raise ValueError('z must be a 1D array.')

  m = core.concrete_or_error(int, m, 'Argument m of lpmn.')
  n = core.concrete_or_error(int, n, 'Argument n of lpmn.')

  if m != n:
    raise NotImplementedError('Computations for m!=n are not yet supported.')

  l_max = n
  is_normalized = False
  p_vals = _gen_associated_legendre(l_max, z, is_normalized)
  p_derivatives = _gen_derivatives(p_vals, z, is_normalized)

  return (p_vals, p_derivatives)


def lpmn_values(m: int, n: int, z: Array, is_normalized: bool) -> Array:
  r"""The associated Legendre functions (ALFs) of the first kind.

  Unlike `lpmn`, this function only computes the values of ALFs.
  The ALFs of the first kind can be used in spherical harmonics. The
  spherical harmonic of degree `l` and order `m` can be written as
  :math:`Y_l^m(\theta, \phi) = N_l^m * P_l^m(\cos \theta) * \exp(i m \phi)`,
  where :math:`N_l^m` is the normalization factor and θ and φ are the
  colatitude and longitude, respectively. :math:`N_l^m` is chosen in the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Set m equal to n and slice the result: lpmn(n, n, z) returns all orders up to n, so extract the m you need from the output.
  2. If slicing the full table is not enough, use scipy.special.lpmn outside of JAX (e.g. precompute constants) or implement the recurrence yourself.
  3. Check the JAX version/release notes; support for m != n may be added later.

Example fix

// before
p, dp = jax.scipy.special.lpmn(2, 5, z)

// after: compute full table with m == n and slice
p, dp = jax.scipy.special.lpmn(5, 5, z)
p_m2 = p[2]  # values for order m=2
Defensive patterns

Strategy: validation

Validate before calling

assert m == n, f'lpmn in JAX requires m == n, got m={m}, n={n}'

Prevention

When it happens

Trigger: Calling jax.scipy.special.lpmn(m, n, z) with integer m != n, e.g. lpmn(2, 5, z). The check happens eagerly after concrete-or-error validation of m and n.

Common situations: Porting SciPy code that uses scipy.special.lpmn(m, n, x) with arbitrary m <= n; computing Legendre functions of a specific order lower than the degree inside JAX pipelines.

Related errors


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