jax-ml/jax · error · TypeError

A ShapeDtypeStruct does not have a value and cannot be used

Error message

A ShapeDtypeStruct does not have a value and cannot be used as a constant in a JAX function.

What it means

ShapeDtypeStruct describes only a shape and dtype — it carries no data — so it cannot be materialized as a constant in an MLIR computation. JAX registers a handler that always raises to catch attempts to use these placeholder objects where actual values are required.

Source

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

      return out
  m = getattr(val, '__jax_array__', None)
  if m is not None:
    return ir_constant(m())
  raise TypeError(f"No constant handler for type: {type(val)}")


def _numpy_array_constant(x: np.ndarray | np.generic) -> ir.Value:
  return hlo.constant(_numpy_array_attribute(x))


def _masked_array_constant_handler(*args, **kwargs):
  raise ValueError("numpy masked arrays are not supported as direct inputs to JAX functions. "
                   "Use arr.filled() to convert the value to a standard numpy array.")

register_constant_handler(np.ma.MaskedArray, _masked_array_constant_handler)

def _shape_dtype_struct_constant_handler(*args, **kwargs):
  raise TypeError("A ShapeDtypeStruct does not have a value and cannot be "
                  "used as a constant in a JAX function.")

register_constant_handler(core.ShapeDtypeStruct,
                          _shape_dtype_struct_constant_handler)

def _ndarray_constant_handler(val: np.ndarray | np.generic,
                              aval: core.AbstractValue | None) -> IrValues:
  """Constant handler for ndarray literals, handling zero-size strides.

  In most cases this function calls _numpy_array_constant(val) except it has
  special handling of arrays with any strides of size zero: for those, it
  generates appropriate calls to NumpyArrayConstant, Broadcast, and Transpose
  to avoid staging in large literals that might arise from np.zeros or np.ones
  or the output of lax.broadcast (which uses np.broadcast_to which in turn
  uses size-zero strides).

  Args:
    val: an ndarray.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Replace the struct with a real array: jnp.zeros(s.shape, s.dtype)
  2. Use the struct only for metadata (shape/dtype) in APIs that accept it (e.g. eval_shape, pjit in_specs)
  3. Check for accidental mixing of metadata dicts and data in argument lists

Example fix

# before
out = jitted_fn(ShapeDtypeStruct((3,), jnp.float32))  # TypeError

# after
out = jitted_fn(jnp.zeros((3,), jnp.float32))
Defensive patterns

Strategy: type-guard

Validate before calling

import jax

def ensure_array(v):
    if isinstance(v, jax.ShapeDtypeStruct):
        return jnp.zeros(v.shape, v.dtype)
    return v

Type guard

import jax

def is_shape_dtype_struct(v) -> bool:
    return isinstance(v, jax.ShapeDtypeStruct)

Prevention

When it happens

Trigger: Passing a jax.ShapeDtypeStruct (e.g. from jax.eval_shape, shard_map args, or checkify/pjit specs) as a value argument to a jitted function instead of using it to describe metadata; feeding aval-shaped placeholders into computations that get traced as constants.

Common situations: Prototyping with eval_shape outputs and then reusing the struct as data; wiring donation/partitioning specs where a struct leaks into argument position; confusing ShapeDtypeStruct with zeros(struct.shape, struct.dtype).

Related errors


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