{"record":{"id":"34a65b9f1b35e8fb","repo":"jax-ml/jax","slug":"unsupported-aval-type-type-v","errorCode":null,"errorMessage":"Unsupported aval type: {type(v)}","messagePattern":"Unsupported aval type: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/checkify.py","lineNumber":990,"sourceCode":"\ndef shard_map_error_check(\n    error: Error, enabled_errors, *vals_in,\n    jaxpr: core.Jaxpr, in_specs, out_specs, **kwargs\n):\n  if (mesh := kwargs.get('mesh')) is None:\n    raise ValueError('Mesh must be provided for shard_map with checkify.')\n\n  err_vals, err_tree = jtu.tree_flatten(error)\n  num_error_vals = len(err_vals)\n  # Replicated sharding for in errors.\n  new_in_specs = (*([P()] * num_error_vals), *in_specs)\n  new_vals_in = [*err_vals, *vals_in]\n  in_avals = list(map(core.typeof, new_vals_in))\n  manual_axes = kwargs.get('newly_manual_axes')\n  check_vma = kwargs.get('check_vma')\n  for i, v in enumerate(in_avals):\n    if not (sharder := core.shard_aval_handlers.get(type(v))):\n      raise ValueError(f'Unsupported aval type: {type(v)}')\n    in_avals[i] = sharder(mesh, manual_axes, check_vma, new_in_specs[i], v)\n\n  with (jshmap._extend_axis_env(mesh, manual_axes),\n        mesh_lib.use_abstract_mesh(jshmap._as_manual_mesh(mesh, manual_axes)),\n        config._check_vma(check_vma)):\n    # jaxpr to checked_jaxpr\n    checked_jaxpr, out_tree, _ = jaxpr_to_checkify_jaxpr(\n        jaxpr, enabled_errors, err_tree, *in_avals\n    )\n  num_out_error_vals = out_tree.num_leaves - len(out_specs)\n\n  def expand_errors_leading_dim(*xs):\n    outs = core.eval_jaxpr(checked_jaxpr, checked_jaxpr.consts, *xs)\n    errs, outs = split_list(outs, [num_out_error_vals])\n    errs = [lax.expand_dims(e, [0]) for e in errs]\n    return *errs, *outs\n\n  with core.extend_axis_env_nd(mesh.shape.items()), config._check_vma(check_vma):","sourceCodeStart":972,"sourceCodeEnd":1008,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/checkify.py#L972-L1008","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove the exotic-typed input from the checkified shard_map region (pass it through outside)","Register a handler in core.shard_aval_handlers for your custom aval type","Update JAX — handlers for newer aval types get added over time"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":"known = set(core.shard_aval_handlers)\nassert all(type(core.typeof(v)) in known for v in jax.tree.leaves(inputs)), 'exotic aval under checkify+shard_map'","typeGuard":"def has_shard_handler(v) -> bool:\n    return type(core.typeof(v)) in core.shard_aval_handlers","tryCatchPattern":null,"preventionTips":["Keep custom aval types out of checkified shard_map scopes","Register shard_aval_handlers when defining custom avals"],"tags":["jax","checkify","shard-map","custom-aval","extensibility"],"backgroundTag":"unsupported-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}