jax-ml/jax · error · NotImplementedError

Expected a ShapeDtypeStruct, but got: {r}

Error message

Expected a ShapeDtypeStruct, but got: {r}

What it means

When computing MLIR result types for a custom Mosaic GPU primitive, every flattened return value must be a ShapeDtypeStruct. Any other object (e.g. a ShapedArray, ref, or FragmentedArray) raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3734

          in_layouts.append(mgpu_layouts.to_layout_attr(t.to_mgpu()))
      case _:
        raise NotImplementedError(
            f"Unsupported aval type: {aval}, {type(aval)}, {t}"
        )
  return in_types, in_layouts, in_transforms


def _custom_primitive_op_results(flat_ret_ty) -> tuple[
    Sequence[ir.Type],
    Sequence[ir.Attribute | None],
]:
  """Returns a tuple containing the list of output MLIR types, and layouts for
  the given JAX return types."""
  results_ty: list[ir.Type] = []
  out_layouts: list[ir.Attribute | None] = []
  for r in flat_ret_ty:
    if not isinstance(r, ShapeDtypeStruct):
      raise NotImplementedError(f"Expected a ShapeDtypeStruct, but got: {r}")
    el_type = mgpu_utils.dtype_to_ir_type(r.dtype)
    if not r.shape:  # scalar case.
      results_ty.append(el_type)
      out_layouts.append(None)
    else:
      results_ty.append(ir.VectorType.get(r.shape, el_type))
      layout = mgpu_layouts.to_layout_attr(r.layout.to_mgpu())
      out_layouts.append(layout)
  return results_ty, out_layouts


def _populate_custom_primitive_op_block(
    ctx: lowering.LoweringRuleContext,
    block: ir.Block,
    mgpu_fn: Callable[..., Any],
    pytree_args,
    in_layouts: Sequence[ir.Attribute],
    in_transforms: Sequence[ir.ArrayAttr],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert return type leaves to ShapeDtypeStruct(shape, dtype) before passing them
2. Audit the out_types of the custom primitive rule so each leaf is a ShapeDtypeStruct

Example fix

// before
out_ty = jax.core.ShapedArray((8, 8), jnp.float32)
// after
out_ty = jax.ShapeDtypeStruct((8, 8), jnp.float32)

assert all(isinstance(r, jax.ShapeDtypeStruct) for r in flat_ret_ty)

def all_sds(rets): return all(isinstance(r, jax.ShapeDtypeStruct) for r in rets)

When it happens

Trigger: Declaring the return type of a custom primitive with leaves that are not jax.ShapeDtypeStruct instances.

Common situations: Mixing ShapedArray (from jax.core) with ShapeDtypeStruct (expected for output specs); custom primitive definitions copied from older examples.

Related errors

• Invalid out_shape type: {type(out_shape)}
• Custom VJP bwd rule {self.bwd} must produce a tuple but got
• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• Group sizes {group_sizes.shape=} must match first dimension
• Explicit sharding is not currently supported in the pallas-t

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