{"record":{"id":"342d38fce39e0ed5","repo":"jax-ml/jax","slug":"expected-a-shapedtypestruct-but-got-r","errorCode":null,"errorMessage":"Expected a ShapeDtypeStruct, but got: {r}","messagePattern":"Expected a ShapeDtypeStruct, but got: (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3734,"sourceCode":"          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.\n      results_ty.append(el_type)\n      out_layouts.append(None)\n    else:\n      results_ty.append(ir.VectorType.get(r.shape, el_type))\n      layout = mgpu_layouts.to_layout_attr(r.layout.to_mgpu())\n      out_layouts.append(layout)\n  return results_ty, out_layouts\n\n\ndef _populate_custom_primitive_op_block(\n    ctx: lowering.LoweringRuleContext,\n    block: ir.Block,\n    mgpu_fn: Callable[..., Any],\n    pytree_args,\n    in_layouts: Sequence[ir.Attribute],\n    in_transforms: Sequence[ir.ArrayAttr],","sourceCodeStart":3716,"sourceCodeEnd":3752,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3716-L3752","documentation":"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.","triggerScenarios":"Declaring the return type of a custom primitive with leaves that are not jax.ShapeDtypeStruct instances.","commonSituations":"Mixing ShapedArray (from jax.core) with ShapeDtypeStruct (expected for output specs); custom primitive definitions copied from older examples.","solutions":["Convert return type leaves to ShapeDtypeStruct(shape, dtype) before passing them","Audit the out_types of the custom primitive rule so each leaf is a ShapeDtypeStruct"],"exampleFix":"// before\nout_ty = jax.core.ShapedArray((8, 8), jnp.float32)\n// after\nout_ty = jax.ShapeDtypeStruct((8, 8), jnp.float32)","handlingStrategy":"type-guard","validationCode":"assert all(isinstance(r, jax.ShapeDtypeStruct) for r in flat_ret_ty)","typeGuard":"def all_sds(rets): return all(isinstance(r, jax.ShapeDtypeStruct) for r in rets)","tryCatchPattern":null,"preventionTips":["Always declare outputs with jax.ShapeDtypeStruct, not ShapedArray"],"tags":["jax","pallas","shapedtypestruct","return-type"],"backgroundTag":"unsupported-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}