{"record":{"id":"009d37d6a9f4cdd4","repo":"jax-ml/jax","slug":"function-dbg-func-src-info-traced-for-dbg-trace","errorCode":null,"errorMessage":"function {dbg.func_src_info} traced for {dbg.traced_for} returned a value of type {type(x)}{extra}, which is not a valid JAX type","messagePattern":"function (.+?) traced for (.+?) returned a value of type (.+?)(.+?), which is not a valid JAX type","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/partial_eval.py","lineNumber":2172,"sourceCode":"    ans = map(dtypes.canonicalize_value, ans)\n    out_tracers = map(partial(trace.to_jaxpr_tracer, source_info=source_info), ans)\n    _check_no_returned_refs(fun.debug_info, out_tracers)\n    jaxpr, consts = trace.frame.to_jaxpr(trace, out_tracers, fun.debug_info,\n                                         source_info)\n    del trace, fun, in_tracers, out_tracers, ans\n  config.enable_checks.value and core.check_jaxpr(jaxpr)\n  return jaxpr, [v.aval for v in jaxpr.outvars], consts\n\ndef _check_returned_jaxtypes(dbg, out_tracers):\n  for i, x in enumerate(out_tracers):\n    try: typeof(x)\n    except TypeError:\n      if (dbg and len(paths := dbg.resolve_result_paths().result_paths) > i and\n          (p := paths[i].removeprefix('result'))):\n        extra = f' at output component {p}'\n      else:\n        extra = ''\n      raise TypeError(\n      f\"function {dbg.func_src_info} traced for {dbg.traced_for} returned a \"\n      f\"value of type {type(x)}{extra}, which is not a valid JAX type\") from None\n\ndef _check_no_returned_refs(\n    dbg: core.DebugInfo,\n    out_tracers: Sequence[DynamicJaxprTracer]\n) -> None:\n  if not config.mutable_array_checks.value: return\n  for i, t in enumerate(out_tracers):\n    a = t.aval\n    if isinstance(a, AbstractRef):\n      result_paths = dbg.resolve_result_paths().safe_result_paths(len(out_tracers))\n      if list(result_paths) == [\"result\"]: result_paths = [\"\"]  # TODO(mattjj): fix in callee\n      loc = result_paths[i] and f' at output tree path {result_paths[i]}'\n      frame = t._trace.frame\n      v = t.val\n      eqns = frame.get_eqns()\n      # TODO(dougalm): something more efficient","sourceCodeStart":2154,"sourceCodeEnd":2190,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/partial_eval.py#L2154-L2190","documentation":"When JAX traces a function (for jit, grad, vmap, etc.), every returned value must be convertible to a JAX type (Array/Tracer of float/int/complex/bool dtypes, etc.). If an output is a Python object, string, dict, custom class, etc., tracing raises TypeError, including the offending output path when debug info can resolve it.","triggerScenarios":"A transformed function returning a non-array: a Python string, None mixed into outputs, a custom class, an unsupported dtype (e.g. object/str arrays), or a numpy array of dtype object. pytrees of arrays are fine; non-array leaves are not.","commonSituations":"Functions returning status strings or metadata alongside arrays; returning None instead of an empty tuple; debug code returning the input dict with an extra flag; use of Python format objects in outputs under jit.","solutions":["Return only arrays (or pytrees of arrays); convert or drop non-array values","Move metadata computation outside the jitted function","Replace None outputs with jnp.zeros((), jnp.float32) placeholders or return tuples without None","For the specific failing output, use the 'at output component' path in the message to locate which return element is bad"],"exampleFix":"# before\n@jax.jit\ndef f(x):\n    return x * 2, 'done'\n\n# after\n@jax.jit\ndef f(x):\n    return x * 2\n# handle status outside jitted code","handlingStrategy":"type-guard","validationCode":"def leaves_are_arrays(tree):\n    return all(\n        isinstance(x, (jax.Array, jnp.ndarray)) or np.isscalar(x)\n        for x in jax.tree.leaves(tree)\n    )\nassert leaves_are_arrays(fn(*args))","typeGuard":"import jax\n\ndef is_valid_output_pytree(tree) -> bool:\n    return all(\n        isinstance(l, (jax.Array,)) or np.isscalar(l) and type(l) is not str\n        for l in jax.tree.leaves(tree)\n    )","tryCatchPattern":"try:\n    jitted = jax.jit(fn); out = jitted(x)\nexcept TypeError as e:\n    if 'not a valid JAX type' in str(e):\n        # message names the offending output component; strip non-array leaves","preventionTips":["Return only arrays/pytrees of arrays from transformed functions","Use the 'at output component' path in the message to find the bad output"],"tags":["jax","tracing","type-error","return-value","jit"],"backgroundTag":"non-array-return-from-traced-function","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}