{"record":{"id":"3ab170634da0567c","repo":"jax-ml/jax","slug":"argument-x-of-type-typ-is-not-a-valid-jax","errorCode":null,"errorMessage":"Argument '{x}' of type '{typ}' is not a valid JAX type","messagePattern":"Argument '(.+?)' of type '(.+?)' is not a valid JAX type","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2018,"sourceCode":"# TODO(phawkins): the return type should be AbstractValue.\ndef typeof(x: Any) -> Any:\n  \"\"\"Return the JAX type (i.e. :class:`AbstractValue`) of the input.\n\n  Raises a ``TypeError`` if ``x`` is not a valid JAX type.\n  \"\"\"\n  typ = type(x)\n  if (aval_fn := pytype_aval_mappings.get(typ)):  # fast path\n    return aval_fn(x)\n  for t in typ.__mro__[1:]:\n    if (aval_fn := pytype_aval_mappings.get(t)):\n      return aval_fn(x)\n  if getattr(x, '__jax_array__', None) is not None:\n    raise ValueError(\n        'Triggering __jax_array__() during abstractification is no longer'\n        ' supported. To avoid this error, either explicitly convert your object'\n        ' using jax.numpy.array(), or register your object as a pytree.'\n    )\n  raise TypeError(f\"Argument '{x}' of type '{typ}' is not a valid JAX type\")\n\ndef is_concrete(x):\n  return to_concrete_value(x) is not None\n\ndef to_concrete_value(x):\n  if isinstance(x, Tracer):\n    return x.to_concrete_value()\n  else:\n    return x\n\ndef concretization_function_error(fun, suggest_astype=False):\n  fname = getattr(fun, \"__name__\", fun)\n  fname_context = f\"The problem arose with the `{fname}` function. \"\n  if suggest_astype:\n    fname_context += (\"If trying to convert the data type of a value, \"\n                      f\"try using `x.astype({fun.__name__})` \"\n                      f\"or `jnp.array(x, {fun.__name__})` instead.\")\n  if fun is bool:","sourceCodeStart":2000,"sourceCodeEnd":2036,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2000-L2036","documentation":"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.","triggerScenarios":"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.","commonSituations":"Accidentally passing hyperparameters (strings, config objects) as traced arguments instead of static ones; passing dict configs or Python enums into jitted functions.","solutions":["Move non-array arguments to static_argnums/static_argnames","Convert numeric data with jnp.array before the call","Register custom containers as pytrees","Verify no None/str/sentinel values are passed through"],"exampleFix":"// before\n@jax.jit\ndef step(params, lr_str): ...\nstep(params, '0.01')\n\n// after\n@jax.jit\ndef step(params, lr: float): ...\nstep(params, 0.01)","handlingStrategy":"type-guard","validationCode":"import jax\ndef valid_jax_arg(x):\n    return hasattr(x, 'dtype') or isinstance(x, (int, float, bool, complex)) or jax.tree_util.all_leaves([x])","typeGuard":"def valid_jax_arg(x) -> bool:\n    return hasattr(x, 'dtype') or isinstance(x, (int, float, bool, complex)) or jax.tree_util.all_leaves([x])","tryCatchPattern":null,"preventionTips":["Use static_argnames for strings/config objects","Convert lists to jnp.array at call sites","Lint traced function signatures to allow only arrays and pytrees"],"tags":["jax","type-error","jit","argument-validation"],"backgroundTag":"jax-invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}