jax-ml/jax · error · ValueError

NumPy arrays with zero strides are not supported as MLIR att

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.

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)

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: {type(val)}
• len() of unsized object
• Sharding rule has {len(rule.operand_mappings)} operands, but
• numpy masked arrays are not supported as direct inputs to JA
• No attribute handler defined for type: {type(val)}

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