jax-ml/jax · error · ValueError

NumPy arrays with zero strides are not supported as MLIR…

Error message

NumPy arrays with zero strides are not supported as MLIR attributes

What it means

A numpy array with a zero stride (e.g. created via np.broadcast_to or repetitive slicing) cannot be embedded directly as a DenseElementsAttr, which requires a standard memory layout the MLIR attribute builder can read. JAX raises this before hitting an opaque MLIR binding error.

Solutions

1. Materialize the array with np.ascontiguousarray(arr) before passing it
2. Use arr.copy() or np.broadcast_to(...).copy() to get normal strides
3. Prefer letting JAX broadcast inside the computation (pass the small array and rely on broadcasting rules)

Example fix

# before
bias = np.broadcast_to(np.float32(0.1), (1024,))
jitted_fn(x, bias)  # ValueError: zero strides

# after
bias = np.ascontiguousarray(bias)
jitted_fn(x, bias)

import numpy as np

def ensure_contiguous(a):
    if isinstance(a, np.ndarray) and 0 in a.strides and a.size > 0:
        return np.ascontiguousarray(a)
    return a

x = ensure_contiguous(x)
jitted_fn(x)

def has_zero_strides(a) -> bool:
    import numpy as np
    return isinstance(a, np.ndarray) and a.size > 0 and 0 in a.strides

When it happens

Trigger: Passing arrays produced by np.broadcast_to, np.lib.stride_tricks.as_strided with 0 strides, or certain reshape/transpose chains where some dimension has stride 0 and size > 0 into a jitted function as a constant.

Common situations: Efficient broadcasting patterns from numpy code reused with JAX; creating constant bias vectors via broadcast_to to save memory; zero-stride arrays coming from np.zeros((n,1,1)) style views after operations.

Related errors

• Unsupported scalar attribute type
• `a` array must be integer typed
• After moving axes to end, leading shape of a must match…
• Arguments to jax.numpy.gcd must be integers.
• Arguments to jax.numpy.lcm must be integers.

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

Appendix: source

Thrown at jax/_src/interpreters/mlir.py:431

def _numpy_scalar_attribute(val: Any) -> ir.Attribute:
  mlir_type = dtype_to_ir_type(val.dtype)
  if isinstance(mlir_type, ir.IntegerType):
    return ir.IntegerAttr.get(mlir_type, int(val))
  elif isinstance(mlir_type, ir.FloatType):
    return ir.FloatAttr.get(mlir_type, val)
  else:
    raise TypeError(f"Unsupported scalar attribute type: {type(val)}")

def _numpy_array_attribute(x: np.ndarray | np.generic) -> ir.Attribute:
  element_type = dtype_to_ir_type(x.dtype)
  shape = x.shape
  x = np.ascontiguousarray(x)
  return ir.DenseElementsAttr.get(x, type=element_type, shape=shape)

def _numpy_array_attribute_handler(val: np.ndarray | np.generic) -> ir.Attribute:
  if 0 in val.strides and val.size > 0:
    raise ValueError(
        "NumPy arrays with zero strides are not supported as MLIR attributes")
  if val.dtype == dtypes.float0:
    val = np.zeros(val.shape, dtype=np.bool_)
  if dtypes.is_weakly_typed_scalar(val) or np.isscalar(val):
    return _numpy_scalar_attribute(val)
  else:
    return _numpy_array_attribute(val)

register_attribute_handler(np.ndarray, _numpy_array_attribute_handler)
register_attribute_handler(hashable_array.HashableArray,
                           lambda x: _numpy_array_attribute_handler(x.val))

for _scalar_type in [np.int8, np.int16, np.int32, np.int64,
                     np.uint8, np.uint16, np.uint32, np.uint64,
                     np.float16, np.float32, np.float64,
                     np.complex64, np.complex128,
                     np.bool_, np.longlong, dtypes.bfloat16]:
  register_attribute_handler(_scalar_type, _numpy_array_attribute_handler)

View on GitHub (pinned to 1e1c6a8fc0)