jax-ml/jax · error · ValueError

Invalid out_shape type: {type(out_shape)}

Error message

Invalid out_shape type: {type(out_shape)}

What it means

While converting a declared out_shape into a JAX abstract value (aval), jax found an object that is neither a ShapeDtypeStruct, a known mapped type, nor a duck-typed object exposing both .shape and .dtype. _convert_out_shape_to_aval therefore rejects it with ValueError. The API requires out_shapes to look like shape/dtype descriptors.

Source

Thrown at jax/_src/pallas/core.py:1706

            shape=out_shape.shape, dtype=out_shape.dtype,
            sharding=jax_core.get_cur_mesh_sharding(),
            manual_axis_type=out_shape.manual_axis_type)
      return jax_core.ShapedArray(
          shape=out_shape.shape, dtype=out_shape.dtype,
          sharding=jax_core.get_cur_mesh_sharding())
    case jax_core.ShapedArray():
      return out_shape
    case MemoryRef():
      return out_shape.get_array_aval()
    case hijax.HiType():
      return out_shape
    case _:
      if type(out_shape) in _out_shape_to_aval_mapping:
        return _out_shape_to_aval_mapping[type(out_shape)](
            out_shape
        )
      if not (hasattr(out_shape, "shape") and hasattr(out_shape, "dtype")):
        raise ValueError(f"Invalid out_shape type: {type(out_shape)}")
      return jax_core.ShapedArray(shape=out_shape.shape, dtype=out_shape.dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the shape in jax.ShapeDtypeStruct(shape, dtype) before passing it as out_shape
2. If using a custom class, give it .shape and .dtype properties so duck-typing succeeds
3. Pass jax.core.ShapedArray directly, which is in the known mapping
4. Check you are not accidentally passing dtype or axis names where out_shape is expected

Example fix

# before
kernel = pallas_kernel(fn, out_shape=(8, 8))  # tuple has no .shape/.dtype

# after
out = jax.ShapeDtypeStruct((8, 8), jnp.float32)
kernel = pallas_kernel(fn, out_shape=out)

def validate_out_shape(o):
    assert hasattr(o, 'shape') and hasattr(o, 'dtype'), f'out_shape {type(o)} must expose .shape and .dtype'

from typing import Any
def is_valid_out_shape(o: Any) -> bool:
    return hasattr(o, 'shape') and hasattr(o, 'dtype')

try:
    aval = convert_out_shape(o)
except ValueError as e:
    if 'Invalid out_shape type' in str(e):
        o = jax.ShapeDtypeStruct(o.shape, o.dtype)  # normalize then retry
        aval = convert_out_shape(o)
    else:
        raise

When it happens

Trigger: Passing an int, tuple, numpy dtype, string, or arbitrary object as out_shape/out_shapes entry to a pallas/shard_map-style API that routes through _convert_out_shape_to_aval (jax/_src/pallas/core.py:1706).

Common situations: Passing a bare shape tuple (8, 8) instead of ShapeDtypeStruct((8,8), dtype); passing a custom container that lacks .shape/.dtype attributes; refactoring where out_shape became a dataclass without those attribute names.

Related errors

• Expected a ShapeDtypeStruct, but got: {r}
• numpy masked arrays are not supported as direct inputs to JA
• Invalid compute type {c_type}. Current supported values are
• numpy masked arrays are not supported as direct inputs to JA
• This gmm kernel only supports either (m, k) x (g, k, n) -> (

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