jax-ml/jax · error · NotImplementedError

Unsupported aval type: {aval}, {type(aval)}, {t}

Error message

Unsupported aval type: {aval}, {type(aval)}, {t}

What it means

While building MLIR input types for a custom Pallas primitive on Mosaic GPU, the aval's type fell into no supported case of the match statement. Only ShapeDtypeStruct-like avals mapping to scalar or VectorType are supported; anything else raises NotImplementedError with the aval, its Python type, and the transform t.

Source

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

  for aval, transformed, t in zip(
      flat_arg_avals, flat_transformed_args, flat_arg_types
  ):
    match aval:
      case state.AbstractRef():
        initial_ty = ir.MemRefType(transformed.type)
        in_types.append(initial_ty)
        if mgpu_utils.is_smem_ref(initial_ty):
          in_transforms.append(_ref_type_to_transforms(t))
      case jax_core.ShapedArray() if isinstance(t, SomeLayout):
        el_type = mgpu_utils.dtype_to_ir_type(aval.dtype)
        if len(aval.shape) == 0:
          in_types.append(el_type)
        else:
          vector_type = ir.VectorType.get(aval.shape, el_type)
          in_types.append(vector_type)
          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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Filter out or handle non-array arguments before the custom primitive
2. Wrap the value so its aval is a ShapedArray/ShapeDtypeStruct
3. Implement a new case in the match statement (contributor fix)

assert all(isinstance(a, (jax.core.ShapedArray, jax.ShapeDtypeStruct)) for a in flat_args)

def is_supported_aval(a): return hasattr(a, 'shape') and hasattr(a, 'dtype')

When it happens

Trigger: Passing an argument whose aval is not a ShapedArray or ShapeDtypeStruct (e.g. a token, a ref with unusual memory space, or an extended aval type) to a custom primitive lowering in the input-types construction.

Common situations: Custom primitives receiving effects tokens or non-array leaves; forwarding internal avals not yet supported by mosaic_gpu lowering.

Related errors

• cannot cast {src} to {dst_type}
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Array type must have a `shape` attribute, but got {type(arra

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