jax-ml/jax · error · ValueError

z must be a 1D array.

Error message

z must be a 1D array.

What it means

jax.scipy.special.lpmn requires the output order m to equal the degree n (it computes the full m x n matrix only for the square case). Passing m != n raises NotImplementedError. m and n must also be concrete Python ints.

Source

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

  Returns:
    A 2-tuple of 3D arrays of shape `(l_max + 1, l_max + 1, len(z))` containing
    the values and derivatives of the associated Legendre functions of the
    first kind. The return type matches the type of `z`.

  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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Request the full square case: lpmn(n, n, z) and slice the result rows/columns to the m you need
2. If m > n, slice the first m rows of the (n+1, n+1) output; keep both arguments equal
3. Ensure m and n are Python ints, and pass them as static to jit (static_argnums or functools.partial)

Example fix

// before
p, dp = jax.scipy.special.lpmn(2, 5, z)  # m != n
// after
p, dp = jax.scipy.special.lpmn(5, 5, z)
p_m2, dp_m2 = p[2], dp[2]  # slice the order you need

assert isinstance(m, int) and isinstance(n, int) and m == n, 'lpmn requires m == n; request (n, n) and slice'

def lpmn_args_ok(m, n):
    return isinstance(m, int) and isinstance(n, int) and m == n

When it happens

Trigger: Calling lpmn(m=2, n=5, z) — any m != n combination raises; traced m/n (e.g., inside jit without static args) instead fails the concrete_or_error check just before.

Common situations: Porting scipy.special.lpmn(m, n, x) calls that request a subset of orders (common in spherical harmonics where only some m are needed); passing m/n from batched computation as arrays rather than static ints.

Related errors

• Negative orders for normalization is not implemented yet.
• Riemann zeta function not implemented; pass q != None to com
• series_order must be non-negative.
• series_order must be <= 30.
• z.dtype={} is not supported, see docstring for supported typ

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