jax-ml/jax · error · TypeError

Unsupported scalar attribute type: {type(val)}

Error message

Unsupported scalar attribute type: {type(val)}

What it means

When converting a numpy scalar into an MLIR attribute, JAX maps the dtype to an MLIR IntegerType or FloatType and builds the corresponding attribute. If dtype_to_ir_type returns something else (complex handled elsewhere, or an exotic/unmapped dtype), the scalar cannot be represented and a TypeError is raised.

Source

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

# Attributes

AttributeHandler = Callable[[Any], ir.Attribute]
_attribute_handlers: dict[type[Any], AttributeHandler] = {}

def register_attribute_handler(type_: type[Any], handler_fun: AttributeHandler):
  _attribute_handlers[type_] = handler_fun

def get_attribute_handler(type_: type[Any]) -> AttributeHandler:
  return _attribute_handlers[type_]

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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the value to a supported dtype before passing it: arr.astype(np.float32) / np.int32 / jnp.bfloat16 via JAX
2. Upgrade JAX so the dtype mapping table covers your type
3. Inspect val.dtype and convert exotic dtypes (datetime/object/string) to numeric or string metadata outside the traced function

Example fix

# before
jitted_fn(np.datetime64('2024-01-01'))

# after
jitted_fn(np.int64(19723))  # convert dates to numeric outside JAX

SUPPORTED = {np.dtype(t) for t in (np.bool_, np.int8, np.int16, np.int32, np.int64,
                                       np.uint8, np.uint16, np.uint32, np.uint64,
                                       np.float16, np.float32, np.float64)}

def cast_to_supported(arr):
    if arr.dtype not in SUPPORTED:
        return arr.astype(np.float32)
    return arr

When it happens

Trigger: Lowering a constant whose numpy scalar dtype maps to a non-integer/non-float MLIR type — typically a dtype that JAX does not support (e.g. float128 on platforms where it is not mapped, datetime64, timedelta64, or object/string scalars reaching this path).

Common situations: Environment-dependent dtype mismatches: np.longdouble on ARM/Windows, datasets containing np.datetime64 columns, older JAX versions lacking mappings for newer dtypes, or bfloat16 conversions done via unsupported numpy paths.

Related errors

• NumPy arrays with zero strides are not supported as MLIR att
• `a` array must be integer typed
• jax.numpy.var does not yet support real dtype parameters whe
• dtype parameter is not supported by Buffer.__array__.
• primal and tangent arguments to jax.jvp do not match; dtypes

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