jax-ml/jax · error · TypeError

Argument '{x}' of type '{typ}' is not a valid JAX type

Error message

Argument '{x}' of type '{typ}' is not a valid JAX type

What it means

TypeError from JAX typeof(): the argument's type has no registered aval mapping, no __jax_array__, and no dtype, so it is not a valid JAX type and cannot be traced.

Source

Thrown at jax/_src/core.py:2018

# TODO(phawkins): the return type should be AbstractValue.
def typeof(x: Any) -> Any:
  """Return the JAX type (i.e. :class:`AbstractValue`) of the input.

  Raises a ``TypeError`` if ``x`` is not a valid JAX type.
  """
  typ = type(x)
  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path
    return aval_fn(x)
  for t in typ.__mro__[1:]:
    if (aval_fn := pytype_aval_mappings.get(t)):
      return aval_fn(x)
  if getattr(x, '__jax_array__', None) is not None:
    raise ValueError(
        'Triggering __jax_array__() during abstractification is no longer'
        ' supported. To avoid this error, either explicitly convert your object'
        ' using jax.numpy.array(), or register your object as a pytree.'
    )
  raise TypeError(f"Argument '{x}' of type '{typ}' is not a valid JAX type")

def is_concrete(x):
  return to_concrete_value(x) is not None

def to_concrete_value(x):
  if isinstance(x, Tracer):
    return x.to_concrete_value()
  else:
    return x

def concretization_function_error(fun, suggest_astype=False):
  fname = getattr(fun, "__name__", fun)
  fname_context = f"The problem arose with the `{fname}` function. "
  if suggest_astype:
    fname_context += ("If trying to convert the data type of a value, "
                      f"try using `x.astype({fun.__name__})` "
                      f"or `jnp.array(x, {fun.__name__})` instead.")
  if fun is bool:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move non-array arguments to static_argnums/static_argnames
2. Convert numeric data with jnp.array before the call
3. Register custom containers as pytrees
4. Verify no None/str/sentinel values are passed through

Example fix

// before
@jax.jit
def step(params, lr_str): ...
step(params, '0.01')

// after
@jax.jit
def step(params, lr: float): ...
step(params, 0.01)

import jax
def valid_jax_arg(x):
    return hasattr(x, 'dtype') or isinstance(x, (int, float, bool, complex)) or jax.tree_util.all_leaves([x])

def valid_jax_arg(x) -> bool:
    return hasattr(x, 'dtype') or isinstance(x, (int, float, bool, complex)) or jax.tree_util.all_leaves([x])

When it happens

Trigger: Passing unsupported Python objects (strings, dicts, generators, arbitrary class instances) as arguments to jitted/grad/vmapped functions, where typeof(x) is called during argument processing.

Common situations: Accidentally passing hyperparameters (strings, config objects) as traced arguments instead of static ones; passing dict configs or Python enums into jitted functions.

Related errors

• `compute_on`'s compute_type argument must be a string.
• function {dbg.func_src_info} traced for {dbg.traced_for} ret
• args argument to jax.scipy.optimize.minimize must be a tuple
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li

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