{"record":{"id":"6115fa5bdf0d5d75","repo":"jax-ml/jax","slug":"unsupported-aval-type-aval-type-aval-t","errorCode":null,"errorMessage":"Unsupported aval type: {aval}, {type(aval)}, {t}","messagePattern":"Unsupported aval type: (.+?), (.+?), (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3718,"sourceCode":"  for aval, transformed, t in zip(\n      flat_arg_avals, flat_transformed_args, flat_arg_types\n  ):\n    match aval:\n      case state.AbstractRef():\n        initial_ty = ir.MemRefType(transformed.type)\n        in_types.append(initial_ty)\n        if mgpu_utils.is_smem_ref(initial_ty):\n          in_transforms.append(_ref_type_to_transforms(t))\n      case jax_core.ShapedArray() if isinstance(t, SomeLayout):\n        el_type = mgpu_utils.dtype_to_ir_type(aval.dtype)\n        if len(aval.shape) == 0:\n          in_types.append(el_type)\n        else:\n          vector_type = ir.VectorType.get(aval.shape, el_type)\n          in_types.append(vector_type)\n          in_layouts.append(mgpu_layouts.to_layout_attr(t.to_mgpu()))\n      case _:\n        raise NotImplementedError(\n            f\"Unsupported aval type: {aval}, {type(aval)}, {t}\"\n        )\n  return in_types, in_layouts, in_transforms\n\n\ndef _custom_primitive_op_results(flat_ret_ty) -> tuple[\n    Sequence[ir.Type],\n    Sequence[ir.Attribute | None],\n]:\n  \"\"\"Returns a tuple containing the list of output MLIR types, and layouts for\n  the given JAX return types.\"\"\"\n  results_ty: list[ir.Type] = []\n  out_layouts: list[ir.Attribute | None] = []\n  for r in flat_ret_ty:\n    if not isinstance(r, ShapeDtypeStruct):\n      raise NotImplementedError(f\"Expected a ShapeDtypeStruct, but got: {r}\")\n    el_type = mgpu_utils.dtype_to_ir_type(r.dtype)\n    if not r.shape:  # scalar case.","sourceCodeStart":3700,"sourceCodeEnd":3736,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3700-L3736","documentation":"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.","triggerScenarios":"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.","commonSituations":"Custom primitives receiving effects tokens or non-array leaves; forwarding internal avals not yet supported by mosaic_gpu lowering.","solutions":["Filter out or handle non-array arguments before the custom primitive","Wrap the value so its aval is a ShapedArray/ShapeDtypeStruct","Implement a new case in the match statement (contributor fix)"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":"assert all(isinstance(a, (jax.core.ShapedArray, jax.ShapeDtypeStruct)) for a in flat_args)","typeGuard":"def is_supported_aval(a): return hasattr(a, 'shape') and hasattr(a, 'dtype')","tryCatchPattern":null,"preventionTips":["Pass only array-like leaves into custom Mosaic primitives"],"tags":["jax","pallas","aval","not-implemented","mlir"],"backgroundTag":"unsupported-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}