jax-ml/jax · error · TypeError

function {dbg.func_src_info} traced for {dbg.traced_for} ret

Error message

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

What it means

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.

Source

Thrown at jax/_src/interpreters/partial_eval.py:2172

    ans = map(dtypes.canonicalize_value, ans)
    out_tracers = map(partial(trace.to_jaxpr_tracer, source_info=source_info), ans)
    _check_no_returned_refs(fun.debug_info, out_tracers)
    jaxpr, consts = trace.frame.to_jaxpr(trace, out_tracers, fun.debug_info,
                                         source_info)
    del trace, fun, in_tracers, out_tracers, ans
  config.enable_checks.value and core.check_jaxpr(jaxpr)
  return jaxpr, [v.aval for v in jaxpr.outvars], consts

def _check_returned_jaxtypes(dbg, out_tracers):
  for i, x in enumerate(out_tracers):
    try: typeof(x)
    except TypeError:
      if (dbg and len(paths := dbg.resolve_result_paths().result_paths) > i and
          (p := paths[i].removeprefix('result'))):
        extra = f' at output component {p}'
      else:
        extra = ''
      raise TypeError(
      f"function {dbg.func_src_info} traced for {dbg.traced_for} returned a "
      f"value of type {type(x)}{extra}, which is not a valid JAX type") from None

def _check_no_returned_refs(
    dbg: core.DebugInfo,
    out_tracers: Sequence[DynamicJaxprTracer]
) -> None:
  if not config.mutable_array_checks.value: return
  for i, t in enumerate(out_tracers):
    a = t.aval
    if isinstance(a, AbstractRef):
      result_paths = dbg.resolve_result_paths().safe_result_paths(len(out_tracers))
      if list(result_paths) == ["result"]: result_paths = [""]  # TODO(mattjj): fix in callee
      loc = result_paths[i] and f' at output tree path {result_paths[i]}'
      frame = t._trace.frame
      v = t.val
      eqns = frame.get_eqns()
      # TODO(dougalm): something more efficient

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Return only arrays (or pytrees of arrays); convert or drop non-array values
  2. Move metadata computation outside the jitted function
  3. Replace None outputs with jnp.zeros((), jnp.float32) placeholders or return tuples without None
  4. For the specific failing output, use the 'at output component' path in the message to locate which return element is bad

Example fix

# before
@jax.jit
def f(x):
    return x * 2, 'done'

# after
@jax.jit
def f(x):
    return x * 2
# handle status outside jitted code
Defensive patterns

Strategy: type-guard

Validate before calling

def leaves_are_arrays(tree):
    return all(
        isinstance(x, (jax.Array, jnp.ndarray)) or np.isscalar(x)
        for x in jax.tree.leaves(tree)
    )
assert leaves_are_arrays(fn(*args))

Type guard

import jax

def is_valid_output_pytree(tree) -> bool:
    return all(
        isinstance(l, (jax.Array,)) or np.isscalar(l) and type(l) is not str
        for l in jax.tree.leaves(tree)
    )

Try / catch

try:
    jitted = jax.jit(fn); out = jitted(x)
except TypeError as e:
    if 'not a valid JAX type' in str(e):
        # message names the offending output component; strip non-array leaves

Prevention

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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