{"record":{"id":"d896f84de536cf87","repo":"jax-ml/jax","slug":"compute-on-s-compute-type-argument-must-be-a-str","errorCode":null,"errorMessage":"`compute_on`'s compute_type argument must be a string.","messagePattern":"`compute_on`'s compute_type argument must be a string\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/compute_on.py","lineNumber":74,"sourceCode":"  if (c_type not in {'device_host', 'device', 'tpu_sparsecore'}\n      and not c_type.startswith(\"gpu_stream:\")):\n    raise ValueError(\n        f'Invalid compute type {c_type}. Current supported values '\n        'are `device_host`, `device`, `tpu_sparsecore`, and `gpu_stream:#`.')\n\n\ndef compute_on(f=None, *, compute_type, out_memory_spaces,\n                compiler_options=None):\n  kwargs = dict(compute_type=compute_type, out_memory_spaces=out_memory_spaces,\n                compiler_options=compiler_options)\n  if f is None:\n    return lambda g: _compute_on(g, **kwargs)\n  return _compute_on(f, **kwargs)\n\n\ndef _compute_on(f, *, compute_type, out_memory_spaces, compiler_options):\n  if not isinstance(compute_type, str):\n    raise TypeError(\"`compute_on`'s compute_type argument must be a string.\")\n  _check_valid(compute_type)\n\n  def wrapped(*args, **kwargs):\n    nonlocal compiler_options\n    dbg = debug_info('compute_on', f, args, kwargs)\n    args_flat, in_tree = tracing_registry.flatten((args, kwargs))\n    in_avals = tuple(core.shaped_abstractify(x) for x in args_flat)\n    with extend_compute_type(compute_type):\n      jaxpr, out_avals = pe.trace_to_jaxpr(\n          f, ft.treedef_args_to_ft(in_tree, in_avals), dbg)\n      out_tree = out_avals.tree\n      if any(isinstance(c, core.Tracer) for c in jaxpr.consts):\n        jaxpr, consts = pe.separate_consts(jaxpr)\n      else:\n        consts = []\n    out_memory_spaces_flat = flatten_axes(\n        \"compute_on out_memory_spaces\", out_tree, out_memory_spaces)\n    if compute_type == 'tpu_sparsecore' and compiler_options is not None:","sourceCodeStart":56,"sourceCodeEnd":92,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/compute_on.py#L56-L92","documentation":"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.","triggerScenarios":"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'.","commonSituations":"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.","solutions":["Pass compute_type as a literal string such as 'cpu', 'gpu', or 'tpu'","If the value comes from a variable, coerce/validate it first: compute_type = str(compute_type) if compute_type else 'cpu'","Check for typos — the keyword is compute_type, not device or backend"],"exampleFix":"// before\nf = jax.compute_on(fn, compute_type=jax.devices()[0])\n// after\nf = jax.compute_on(fn, compute_type='gpu')","handlingStrategy":"type-guard","validationCode":"def is_valid_compute_type(ct):\n    return isinstance(ct, str) and ct in ('cpu', 'gpu', 'tpu')\n\nct = ct or 'cpu'\nassert is_valid_compute_type(ct), f'bad compute_type: {ct!r}'","typeGuard":"def is_compute_type_str(ct: object) -> TypeGuard[str]:\n    return isinstance(ct, str)","tryCatchPattern":null,"preventionTips":["Always pass compute_type as a lowercase string literal","Default the variable explicitly: compute_type = compute_type or 'cpu'","Validate backend strings at the edge of your config loading"],"tags":["jax","type-error","compute-on","argument-validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}