jax-ml/jax · error · ValueError

negative dimensions are not allowed, got {N} and {M}

Error message

negative dimensions are not allowed, got {N} and {M}

What it means

jnp.eye validates that the requested number of rows (N) and columns (M) are non-negative after canonicalization to integer dimensions. A negative N or M is meaningless for an eye/identity-like matrix and raises this ValueError echoing both values.

Source

Thrown at jax/_src/numpy/lax_numpy.py:5810

    return api.device_put(output, device=device)
  return output


def _eye(N: DimSize, M: DimSize | None = None,
        k: int | ArrayLike = 0,
        dtype: DTypeLike | None = None) -> Array:
  dtype = dtypes.check_and_canonicalize_user_dtype(
      float if dtype is None else dtype, "eye")
  if isinstance(k, int):
    k = lax._clip_int_to_valid_range(k, np.int32,
                                              "`argument `k` of jax.numpy.eye")
  offset = util.ensure_arraylike("eye", k)
  if not (offset.shape == () and dtypes.issubdtype(offset.dtype, np.integer)):
    raise ValueError(f"k must be a scalar integer; got {k}")
  N_int = core.canonicalize_dim(N, "argument of 'N' jnp.eye()")
  M_int = N_int if M is None else core.canonicalize_dim(M, "argument 'M' of jnp.eye()")
  if N_int < 0 or M_int < 0:
    raise ValueError(f"negative dimensions are not allowed, got {N} and {M}")
  i = lax.broadcasted_iota(offset.dtype, (N_int, M_int), 0)
  j = lax.broadcasted_iota(offset.dtype, (N_int, M_int), 1)
  return (i + offset == j).astype(dtype)


@export
def identity(n: DimSize, dtype: DTypeLike | None = None) -> Array:
  """Create a square identity matrix

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

  Args:
    n: integer specifying the size of each array dimension.
    dtype: optional dtype; defaults to floating point.

  Returns:
    Identity array of shape ``(n, n)``.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp computed dimensions: N = max(N, 0) before calling jnp.eye
2. Fix the upstream size arithmetic (e.g. reduce padding, validate input lengths)
3. Add an assert on expected minimum size before building the matrix

Example fix

// before
eye = jnp.eye(x.shape[0] - y.shape[0])  # can be negative
// after
n = max(x.shape[0] - y.shape[0], 0)
eye = jnp.eye(n)

N, M = max(int(N), 0), (max(int(M), 0) if M is not None else None)
eye = jnp.eye(N, M)

When it happens

Trigger: jnp.eye(-3), jnp.eye(3, -1), or N/M derived from shapes that can be negative, e.g. N = a.shape[0] - b.shape[0] where b is longer than a; padding computations like N = n - 2*pad with pad too large.

Common situations: Padding/cropping arithmetic that overshoots; shape propagation in variable-length pipelines where the computed size goes negative on short inputs.

Related errors

• k must be a scalar integer; got {k}
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va
• Invalid value "{default_str}" for JAX flag {name}

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