{"record":{"id":"beaf77d527d30584","repo":"jax-ml/jax","slug":"can-t-convert-to-ref-aval","errorCode":null,"errorMessage":"Can't convert to ref: {aval}","messagePattern":"Can't convert to ref: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":3652,"sourceCode":"            f\" collective_axes={ctx.module_ctx.axis_names.wg} to your\"\n            \" run_scoped if you intend all threads to share the same\"\n            f\" allocation (currently collective_axes={collective_axes}).\"\n        )\n      if isinstance(aval.dtype, gpu_core.BarrierType):\n        barrier = _get_barrier(aval, ctx.estimator_ctx.arrival_multiplier)\n        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)\n        input_refs.append(alloc_stack.enter_context(barrier_ctx))\n        should_discharge.append(False)\n        continue\n      if isinstance(aval.dtype, gpu_core.ClusterBarrierType):\n        barrier = _get_cluster_barrier(aval, ctx.module_ctx.axis_names)\n        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)\n        input_refs.append(alloc_stack.enter_context(barrier_ctx))\n        should_discharge.append(False)\n        continue\n\n      if not isinstance(aval, state_types.AbstractRef):\n        raise ValueError(f\"Can't convert to ref: {aval}\")\n      if aval.memory_space == gpu_core.SMEM:\n        input_ref = alloc_stack.enter_context(\n            ctx.module_ctx.scratch_view(\n                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype)\n            )\n        )\n        input_refs.append(input_ref)\n        should_discharge.append(False)\n      elif aval.memory_space == gpu_core.TMEM:\n        input_ref = alloc_stack.enter_context(\n            ctx.module_ctx.alloc_tmem(\n                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype),\n                layout=aval.layout,  # pyrefly: ignore[missing-attribute]\n            )\n        )\n        input_refs.append(input_ref)\n        should_discharge.append(False)\n","sourceCodeStart":3634,"sourceCodeEnd":3670,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L3634-L3670","documentation":"run_scoped can only manage allocations for known ref types (WGMMA accumulators, barriers, AbstractRef). If an invar of the scoped jaxpr has an aval that is not state_types.AbstractRef (and not one of the special cases), lowering raises ValueError.","triggerScenarios":"Passing a plain value or a non-ref abstract value as an invar to the run_scoped body, e.g. a Future/token or a custom abstract type that isn't a Ref.","commonSituations":"Tracing bugs where a non-Ref leaks into run_scoped invars; upgrading JAX versions where new aval kinds appear in scoped jaxprs before Mosaic supports them.","solutions":["Ensure everything passed to the run_scoped body is a proper pl state Ref","Wrap raw values in refs (e.g. via pl.run_scoped semantics) before the call","If hitting after a JAX upgrade, pin to a compatible jax version and report upstream"],"exampleFix":"// before\npl.run_scoped(lambda s: body(s), x)  # x is a plain array\n// after\npl.run_scoped(lambda s: body(s), x_ref)  # x_ref = pl.ref(x)","handlingStrategy":"type-guard","validationCode":"assert all(isinstance(v.aval, state_types.AbstractRef) or is_special(v.aval) for v in scoped_jaxpr.invars)","typeGuard":"def is_scoped_ref(v) -> bool:\n    import jax._src.pallas.state as state_types\n    return isinstance(v.aval, state_types.AbstractRef)","tryCatchPattern":null,"preventionTips":["Only pass Refs to run_scoped bodies","Wrap arrays in refs explicitly","Pin compatible jax versions"],"tags":["jax","pallas","run-scoped","ref","type-error"],"backgroundTag":"invalid-ref-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}