{"record":{"id":"94e2e709c0baf1d3","repo":"jax-ml/jax","slug":"unsupported-type-for-json-serialization-type-obj","errorCode":null,"errorMessage":"Unsupported type for JSON serialization: {type(obj)} ({obj})","messagePattern":"Unsupported type for JSON serialization: (.+?) \\((.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":3074,"sourceCode":"    target_name, ctx, x, cfg, called_computations=None\n):\n  out_aval, = ctx.avals_out\n  future_type = mlir.aval_to_ir_type(ctx.module_context, out_aval.inner_aval)\n\n  cfg = dict(cfg)\n  if \"channel_handle\" in cfg:\n    cfg[\"channel_id\"] = cfg.pop(\"channel_handle\").handle\n  if \"use_global_device_ids\" in cfg:\n    cfg[\"use_global_device_ids\"] = cfg[\"use_global_device_ids\"].value\n\n  def _json_default(obj):\n    if isinstance(obj, np.integer):\n      return int(obj)\n    if isinstance(obj, np.floating):\n      return float(obj)\n    if isinstance(obj, np.ndarray):\n      return obj.tolist()\n    raise ValueError(\n        f\"Unsupported type for JSON serialization: {type(obj)} ({obj})\"\n    )\n\n  config_str = json.dumps(cfg, default=_json_default)\n  frontend_attrs = mlir.ir_attribute({\"async_collective_config\": config_str})\n\n  return mlir.custom_call(\n      call_target_name=target_name,\n      result_types=[future_type],\n      operands=[x],\n      extra_attributes={\"mhlo.frontend_attributes\": frontend_attrs},\n      api_version=1,\n      called_computations=[c.name.value for c in called_computations or []],\n  ).results\n\n\ndef _async_done_lowering(target_name, ctx, x):\n  out_aval, = ctx.avals_out","sourceCodeStart":3056,"sourceCodeEnd":3092,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L3056-L3092","documentation":"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.","triggerScenarios":"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.","commonSituations":"Extending JAX or building configs for async collectives with rich Python objects; version upgrades where a config field changed from int to an object.","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"],"exampleFix":"# before\ncfg = {'iters': np.int32(3), 'sched': MySchedule(2)}\n# after\ncfg = {'iters': int(np.int32(3)), 'sched': {'step': 2}}","handlingStrategy":"validation","validationCode":"def sanitize(cfg):\n    return {k: (int(v) if isinstance(v, np.integer) else\n               float(v) if isinstance(v, np.floating) else\n               v.tolist() if isinstance(v, np.ndarray) else v)\n            for k, v in cfg.items()}","typeGuard":"def is_json_safe(obj) -> bool:\n    import json\n    try:\n        json.dumps(obj, default=lambda o: None)\n        return True\n    except TypeError:\n        return False","tryCatchPattern":"catch ValueError from json.dumps and convert offending values to primitives, retry once","preventionTips":["Build collective configs from plain Python primitives only"],"tags":["jax","json-serialization","async-collectives","config"],"backgroundTag":"json-serialization-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}