jax-ml/jax · error · TypeError

`compute_on`'s compute_type argument must be a string.

Error message

`compute_on`'s compute_type argument must be a string.

What it means

jax.compute_on requires its compute_type keyword argument to be a Python string naming a hardware/compiler backend (e.g. 'cpu', 'gpu', 'tpu'). The decorated function wrapper validates the type eagerly at decoration time and raises TypeError for any non-string value (None, an Aval, a device object, etc.). This catches misconfiguration before tracing begins.

Source

Thrown at jax/_src/compute_on.py:74

  if (c_type not in {'device_host', 'device', 'tpu_sparsecore'}
      and not c_type.startswith("gpu_stream:")):
    raise ValueError(
        f'Invalid compute type {c_type}. Current supported values '
        'are `device_host`, `device`, `tpu_sparsecore`, and `gpu_stream:#`.')


def compute_on(f=None, *, compute_type, out_memory_spaces,
                compiler_options=None):
  kwargs = dict(compute_type=compute_type, out_memory_spaces=out_memory_spaces,
                compiler_options=compiler_options)
  if f is None:
    return lambda g: _compute_on(g, **kwargs)
  return _compute_on(f, **kwargs)


def _compute_on(f, *, compute_type, out_memory_spaces, compiler_options):
  if not isinstance(compute_type, str):
    raise TypeError("`compute_on`'s compute_type argument must be a string.")
  _check_valid(compute_type)

  def wrapped(*args, **kwargs):
    nonlocal compiler_options
    dbg = debug_info('compute_on', f, args, kwargs)
    args_flat, in_tree = tracing_registry.flatten((args, kwargs))
    in_avals = tuple(core.shaped_abstractify(x) for x in args_flat)
    with extend_compute_type(compute_type):
      jaxpr, out_avals = pe.trace_to_jaxpr(
          f, ft.treedef_args_to_ft(in_tree, in_avals), dbg)
      out_tree = out_avals.tree
      if any(isinstance(c, core.Tracer) for c in jaxpr.consts):
        jaxpr, consts = pe.separate_consts(jaxpr)
      else:
        consts = []
    out_memory_spaces_flat = flatten_axes(
        "compute_on out_memory_spaces", out_tree, out_memory_spaces)
    if compute_type == 'tpu_sparsecore' and compiler_options is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass compute_type as a literal string such as 'cpu', 'gpu', or 'tpu'
2. If the value comes from a variable, coerce/validate it first: compute_type = str(compute_type) if compute_type else 'cpu'
3. Check for typos — the keyword is compute_type, not device or backend

Example fix

// before
f = jax.compute_on(fn, compute_type=jax.devices()[0])
// after
f = jax.compute_on(fn, compute_type='gpu')

def is_valid_compute_type(ct):
    return isinstance(ct, str) and ct in ('cpu', 'gpu', 'tpu')

ct = ct or 'cpu'
assert is_valid_compute_type(ct), f'bad compute_type: {ct!r}'

def is_compute_type_str(ct: object) -> TypeGuard[str]:
    return isinstance(ct, str)

When it happens

Trigger: Calling compute_on(f, compute_type=<non-string>) — e.g. compute_type=None (forgot to pass it), compute_type=jax.devices()[0], or compute_type=some enum/Avax object instead of a string like 'gpu'.

Common situations: Passing a Device instance or an uppercased/env-derived value that ended up None; copying example code that used a variable for the backend type that was never defined.

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

• Argument '{x}' of type '{typ}' is not a valid JAX type
• 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
• {prim_name} takes a scalar pred as argument, got {pred}

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