jax-ml/jax · error · ValueError

Unsupported aval type: {type(v)}

Error message

Unsupported aval type: {type(v)}

What it means

While re-sharding checkify's error values for shard_map, an input abstract value's type has no registered core.shard_aval_handlers entry, so checkify cannot compute its sharded aval.

Source

Thrown at jax/_src/checkify.py:990

def shard_map_error_check(
    error: Error, enabled_errors, *vals_in,
    jaxpr: core.Jaxpr, in_specs, out_specs, **kwargs
):
  if (mesh := kwargs.get('mesh')) is None:
    raise ValueError('Mesh must be provided for shard_map with checkify.')

  err_vals, err_tree = jtu.tree_flatten(error)
  num_error_vals = len(err_vals)
  # Replicated sharding for in errors.
  new_in_specs = (*([P()] * num_error_vals), *in_specs)
  new_vals_in = [*err_vals, *vals_in]
  in_avals = list(map(core.typeof, new_vals_in))
  manual_axes = kwargs.get('newly_manual_axes')
  check_vma = kwargs.get('check_vma')
  for i, v in enumerate(in_avals):
    if not (sharder := core.shard_aval_handlers.get(type(v))):
      raise ValueError(f'Unsupported aval type: {type(v)}')
    in_avals[i] = sharder(mesh, manual_axes, check_vma, new_in_specs[i], v)

  with (jshmap._extend_axis_env(mesh, manual_axes),
        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, manual_axes)),
        config._check_vma(check_vma)):
    # jaxpr to checked_jaxpr
    checked_jaxpr, out_tree, _ = jaxpr_to_checkify_jaxpr(
        jaxpr, enabled_errors, err_tree, *in_avals
    )
  num_out_error_vals = out_tree.num_leaves - len(out_specs)

  def expand_errors_leading_dim(*xs):
    outs = core.eval_jaxpr(checked_jaxpr, checked_jaxpr.consts, *xs)
    errs, outs = split_list(outs, [num_out_error_vals])
    errs = [lax.expand_dims(e, [0]) for e in errs]
    return *errs, *outs

  with core.extend_axis_env_nd(mesh.shape.items()), config._check_vma(check_vma):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the exotic-typed input from the checkified shard_map region (pass it through outside)
2. Register a handler in core.shard_aval_handlers for your custom aval type
3. Update JAX — handlers for newer aval types get added over time

known = set(core.shard_aval_handlers)
assert all(type(core.typeof(v)) in known for v in jax.tree.leaves(inputs)), 'exotic aval under checkify+shard_map'

def has_shard_handler(v) -> bool:
    return type(core.typeof(v)) in core.shard_aval_handlers

When it happens

Trigger: checkify composed with shard_map where an input is an exotic aval type (not standard ShapedArray/Token), e.g. custom trial or extension types from a third-party custom-aval library.

Common situations: Custom JAX extensions (new aval types) used inside checkified shard_map code; internal/nightly JAX where a new aval type lacks a sharding handler yet.

Related errors

• Mesh must be provided for shard_map with checkify.
• {str(exc)}
• Checkify does not support batched while-loops (checkify-of-v
• {prim_name} takes a scalar pred as argument, got {pred}
• Formatting arguments to checkify.check need to be PyTrees of

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