jax-ml/jax · error · TypeError

Argument '{arg}' of type {type(arg)} is not a valid JAX type

Error message

Argument '{arg}' of type {type(arg)} is not a valid JAX type.

What it means

JAX validates every argument passed into a transform or dispatch path with core.valid_jaxtype; anything that is not a JAX-compatible type (array, tracer, standard scalar container) is rejected with this TypeError. It is raised by check_arg, which guards entry points like grad/jacfwd/pjit dispatch. Its purpose is to fail fast before tracing so users get a clear message instead of a cryptic tracer error.

Source

Thrown at jax/_src/dispatch.py:287

    elif eqn.primitive is shard_map.shard_map_p:
      mesh = eqn.params['mesh']
      if isinstance(mesh, AbstractMesh):
        continue
      source_info = SourceInfo(eqn.source_info, eqn.primitive.name)
      out.extend((NamedSharding(mesh, spec), source_info)
                 for spec in [*eqn.params['in_specs'], *eqn.params['out_specs']])
    elif eqn.primitive is device_put_p:
      source_info = SourceInfo(eqn.source_info, eqn.primitive.name)
      out.extend((s, source_info) for s in eqn.params['devices']
                 if isinstance(s, Sharding) and s.memory_kind is not None)
  for subjaxpr in core.subjaxprs(jaxpr):
    out.extend(get_intermediate_shardings(subjaxpr))
  return out


def check_arg(arg: Any):
  if not core.valid_jaxtype(arg):
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid "
                    "JAX type.")


def needs_check_special() -> bool:
  return config.debug_infs.value or config.debug_nans.value

def check_special(name: str, bufs: Sequence[basearray.Array]) -> None:
  if needs_check_special():
    for buf in bufs:
      _check_special(name, buf.dtype, buf)


def check_special_array(name: str, arr: array.ArrayImpl) -> array.ArrayImpl:
  if needs_check_special():
    if dtypes.issubdtype(arr.dtype, np.inexact):
      for buf in arr._arrays:
        _check_special(name, buf.dtype, buf)
  return arr

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert the offending argument with jnp.asarray(x) (and a valid dtype) before passing it in
If using foreign tensors, convert via numpy: jnp.asarray(x.numpy()) or jnp.from_dlpack(x)
Check for object dtype: np.asarray(x).dtype == object and fix the data source
Inspect the traceback to identify which argument named in the message is invalid

Example fix

// before
loss = jax.grad(model)(raw_python_list, params)
// after
loss = jax.grad(model)(jnp.asarray(raw_python_list, dtype=jnp.float32), params)

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
def is_jax_type(x):
    try:
        jax.core.typeof(x); return True
    except TypeError:
        return False
vals = [v for v in args if not is_jax_type(v)]
assert not vals, f'non-JAX args: {vals}'

Type guard

def is_jax_arg(x) -> bool:
    import jax.numpy as jnp
    return hasattr(x, 'dtype') and hasattr(x, 'shape') or isinstance(x, (int, float, complex, bool))

Try / catch

try:
    result = jax.grad(f)(*args)
except TypeError as e:
    if 'not a valid JAX type' in str(e):
        args = tuple(jnp.asarray(a) if not hasattr(a, 'dtype') else a for a in args)

Prevention

Always convert inputs with jnp.asarray before passing to JAX transforms
Validate pytree leaves with jax.tree_util.tree_map(lambda l: jnp.asarray(l), data)
Avoid object-dtype numpy arrays; cast with astype(float)

When it happens

Trigger: Passing a non-JAX object to jax.grad, jax.jacfwd, jit-compiled/pjit functions, or other transforms: e.g. a numpy object-dtype array, a Python custom class, a torch tensor, None, or a string used as a leaf in a pytree.

Common situations: Passing PyTorch tensors or raw Python objects into JAX functions; object-dtype numpy arrays from pandas; dictionaries with non-array leaves; strings/None accidentally used as data.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

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