jax-ml/jax · error · ValueError

Can't convert to ref: {aval}

Error message

Can't convert to ref: {aval}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:3652

            f" collective_axes={ctx.module_ctx.axis_names.wg} to your"
            " run_scoped if you intend all threads to share the same"
            f" allocation (currently collective_axes={collective_axes})."
        )
      if isinstance(aval.dtype, gpu_core.BarrierType):
        barrier = _get_barrier(aval, ctx.estimator_ctx.arrival_multiplier)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue
      if isinstance(aval.dtype, gpu_core.ClusterBarrierType):
        barrier = _get_cluster_barrier(aval, ctx.module_ctx.axis_names)
        barrier_ctx = ctx.module_ctx.reserve_barrier(barrier)
        input_refs.append(alloc_stack.enter_context(barrier_ctx))
        should_discharge.append(False)
        continue

      if not isinstance(aval, state_types.AbstractRef):
        raise ValueError(f"Can't convert to ref: {aval}")
      if aval.memory_space == gpu_core.SMEM:
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.scratch_view(
                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype)
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)
      elif aval.memory_space == gpu_core.TMEM:
        input_ref = alloc_stack.enter_context(
            ctx.module_ctx.alloc_tmem(
                jax.ShapeDtypeStruct(shape=aval.shape, dtype=aval.dtype),
                layout=aval.layout,  # pyrefly: ignore[missing-attribute]
            )
        )
        input_refs.append(input_ref)
        should_discharge.append(False)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure everything passed to the run_scoped body is a proper pl state Ref
2. Wrap raw values in refs (e.g. via pl.run_scoped semantics) before the call
3. If hitting after a JAX upgrade, pin to a compatible jax version and report upstream

Example fix

// before
pl.run_scoped(lambda s: body(s), x)  # x is a plain array
// after
pl.run_scoped(lambda s: body(s), x_ref)  # x_ref = pl.ref(x)

assert all(isinstance(v.aval, state_types.AbstractRef) or is_special(v.aval) for v in scoped_jaxpr.invars)

def is_scoped_ref(v) -> bool:
    import jax._src.pallas.state as state_types
    return isinstance(v.aval, state_types.AbstractRef)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Unsupported block dim type: {type(b)}
• load_staged_rhs must be an integer or None.
• ref must be a reference
• Unsupported type: {x}
• Can only store to references (got {x_ref}).

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