jax-ml/jax · error · ValueError

k must be a scalar or a rank-1 array of length 1 or m.

Error message

k must be a scalar or a rank-1 array of length 1 or m.

What it means

The integration constants k in jnp.polyint must be broadcastable to exactly m values (one constant per integration step): a scalar, a single-element array, or an array of length exactly m.

Source

Thrown at jax/_src/numpy/polynomial.py:570

    ``m`` elements. The second order integration of the polynomial
    :math:`12 x^2 + 12 x + 6` with the constants ``k=[4, 5]`` is
    :math:`x^4 + 2 x^3 + 3 x^2 + 4 x + 5`:

    >>> jnp.polyint(p, m=2, k=jnp.array([4, 5]))
    Array([1., 2., 3., 4., 5.], dtype=float32)
  """
  m = core.concrete_or_error(operator.index, m, "'m' argument of jnp.polyint")
  k = 0 if k is None else k
  p, k = ensure_arraylike("polyint", p, k)
  p_arr, k_arr = promote_dtypes_inexact(p, k)
  del p, k
  if m < 0:
    raise ValueError("Order of integral must be positive (see polyder)")
  k_arr = atleast_1d(k_arr)
  if len(k_arr) == 1:
    k_arr = full((m,), k_arr[0])
  if k_arr.shape != (m,):
    raise ValueError("k must be a scalar or a rank-1 array of length 1 or m.")
  if m == 0:
    return p_arr
  else:
    grid = (arange(len(p_arr) + m, dtype=p_arr.dtype)[np.newaxis]
            - arange(m, dtype=p_arr.dtype)[:, np.newaxis])
    coeff = maximum(1, grid).prod(0)[::-1]
    return true_divide(concatenate((p_arr, k_arr)), coeff)


@export
@api.jit(static_argnames=('m',))
def polyder(p: ArrayLike, m: int = 1) -> Array:
  r"""Returns the coefficients of the derivative of specified order of a polynomial.

  JAX implementation of :func:`numpy.polyder`.

  Args:
    p: Array of polynomials coefficients.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a scalar k (broadcast to all orders)
2. Make len(k) == m exactly
3. Recompute k when m changes

Example fix

// before
jnp.polyint(p, m=3, k=jnp.array([1.0, 2.0]))
// after
jnp.polyint(p, m=3, k=jnp.array([1.0, 2.0, 3.0]))

import jax.numpy as jnp
k_arr = jnp.atleast_1d(jnp.asarray(k))
if k_arr.size not in (1, m):
    k_arr = jnp.full((m,), k_arr.ravel()[0])  # or raise
result = jnp.polyint(p, m=m, k=k_arr)

When it happens

Trigger: Calling jnp.polyint(p, m=2, k=jnp.array([1.0])) is fine, but k=jnp.array([1.0, 2.0, 3.0]) with m=2 fails; also k of shape (m+1,) or any length != 1 and != m.

Common situations: Hardcoding a list of constants then changing m; passing per-order constants array computed for a different integration order; refactoring numpy code where m changed.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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