jax-ml/jax · error · ValueError

Unsupported type for JSON serialization: {type(obj)} ({obj})

Error message

Unsupported type for JSON serialization: {type(obj)} ({obj})

What it means

When emitting an async collective with a config dict, JAX serializes the config to JSON with a `_json_default` that only handles numpy scalars/arrays plus JSON-native types. Any other object type (e.g. a torch tensor, custom class, or Python object) raises this ValueError.

Source

Thrown at jax/_src/lax/parallel.py:3074

    target_name, ctx, x, cfg, called_computations=None
):
  out_aval, = ctx.avals_out
  future_type = mlir.aval_to_ir_type(ctx.module_context, out_aval.inner_aval)

  cfg = dict(cfg)
  if "channel_handle" in cfg:
    cfg["channel_id"] = cfg.pop("channel_handle").handle
  if "use_global_device_ids" in cfg:
    cfg["use_global_device_ids"] = cfg["use_global_device_ids"].value

  def _json_default(obj):
    if isinstance(obj, np.integer):
      return int(obj)
    if isinstance(obj, np.floating):
      return float(obj)
    if isinstance(obj, np.ndarray):
      return obj.tolist()
    raise ValueError(
        f"Unsupported type for JSON serialization: {type(obj)} ({obj})"
    )

  config_str = json.dumps(cfg, default=_json_default)
  frontend_attrs = mlir.ir_attribute({"async_collective_config": config_str})

  return mlir.custom_call(
      call_target_name=target_name,
      result_types=[future_type],
      operands=[x],
      extra_attributes={"mhlo.frontend_attributes": frontend_attrs},
      api_version=1,
      called_computations=[c.name.value for c in called_computations or []],
  ).results


def _async_done_lowering(target_name, ctx, x):
  out_aval, = ctx.avals_out

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert config values to plain Python primitives (int/float/str/list/dict) before passing
Use `.item()`/`int()`/`float()` on tensor-like values
If you control the code path, extend `_json_default` upstream via a PR or pre-serialize yourself

Example fix

# before
cfg = {'iters': np.int32(3), 'sched': MySchedule(2)}
# after
cfg = {'iters': int(np.int32(3)), 'sched': {'step': 2}}

Defensive patterns

Strategy: validation

Validate before calling

def sanitize(cfg):
    return {k: (int(v) if isinstance(v, np.integer) else
               float(v) if isinstance(v, np.floating) else
               v.tolist() if isinstance(v, np.ndarray) else v)
            for k, v in cfg.items()}

Type guard

def is_json_safe(obj) -> bool:
    import json
    try:
        json.dumps(obj, default=lambda o: None)
        return True
    except TypeError:
        return False

Try / catch

catch ValueError from json.dumps and convert offending values to primitives, retry once

Prevention

Build collective configs from plain Python primitives only

When it happens

Trigger: Passing a configuration object to an async collective API (e.g. async psum/all_gather with a config dict) containing a non-numpy, non-primitive value such as a custom enum or class instance.

Common situations: Extending JAX or building configs for async collectives with rich Python objects; version upgrades where a config field changed from int to an object.

Related errors

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