jax-ml/jax · error · ValueError

k must be a scalar integer; got {k}

Error message

k must be a scalar integer; got {k}

What it means

The k (diagonal offset) argument of jnp.eye must be a scalar with integer dtype after conversion to a JAX array. The check `offset.shape == () and issubdtype(offset.dtype, np.integer)` fails for non-scalar k (arrays with ndim > 0) or non-integer k (float, complex), raising this ValueError.

Source

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

  # TODO(vfdev-5): optimize putting the array directly on the device specified
  # instead of putting it on default device and then on the specific device
  output = _eye(N, M=M, k=k, dtype=dtype)
  if device is not None:
    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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use integer division: k = n // 2 not n / 2
2. Wrap k: jnp.eye(N, k=int(k)) to force a Python int scalar
3. Ensure k is a Python int or 0-d integer array before the call

Example fix

// before
eye = jnp.eye(n, k=n/2)   # float
// after
eye = jnp.eye(n, k=n//2)  # int

k = int(k)  # forces scalar int; raises early if not convertible
eye = jnp.eye(N, k=k)

def is_scalar_int(k) -> bool:
    k = jnp.asarray(k)
    return k.shape == () and jnp.issubdtype(k.dtype, jnp.integer)

When it happens

Trigger: jnp.eye(3, k=1.0) (float offset), jnp.eye(3, k=jnp.array([0])) (1-element array, not scalar), or k computed as float from division e.g. k=n/2.

Common situations: Computing the diagonal offset arithmetically (n//2 vs n/2 bug); passing a traced or batched value for k; passing k as a 0-d float array from external configuration.

Related errors

• cond_fun must return a boolean scalar, but got output type(s
• Indexer must have integer or boolean type, got indexer with
• `a` array must be integer typed
• negative dimensions are not allowed, got {N} and {M}
• Only arrays with 32-bit element types can be converted to sc

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