jax-ml/jax · error · ValueError

Invalid compute type {c_type}. Current supported values are

Error message

Invalid compute type {c_type}. Current supported values are `device_host`, `device` and `tpu_sparsecore`

What it means

JAX maps a user-facing 'compute type' string to an internal string when setting up computations. Only `device_host`, `device`, and `tpu_sparsecore` are recognized; anything else reaches this fallthrough and raises. It is a strict input-validation error on a string enum.

Source

Thrown at jax/_src/interpreters/mlir.py:2819

  tokens_in = ctx.tokens_in.subset(effects)
  out_nodes, tokens = call_lowering(
      name, call_jaxpr, backend, ctx.module_context,
      ctx.avals_in, ctx.avals_out, tokens_in, *args,
      dim_var_values=ctx.dim_var_values,
      const_lowering=ctx.const_lowering)
  ctx.set_tokens_out(ctx.tokens_in.update_tokens(tokens))
  return [lower_with_sharding_in_types(ctx, o, a)
          for o, a in zip(out_nodes, ctx.avals_out)]


def map_compute_type(c_type: str) -> str:
  if c_type == "device_host":
    return "host"
  elif c_type == "device":
    return "dense"
  elif c_type == "tpu_sparsecore":
    return "sparseoffload"
  raise ValueError(f"Invalid compute type {c_type}. Current supported values "
                   "are `device_host`, `device` and `tpu_sparsecore`")


def _update_frontend_attributes(op, attrs):
  if isinstance(op, ir.Block):
    return
  if attr_array := op.attributes.get("mhlo.frontend_attributes"):
    assert isinstance(attr_array, ir.DictAttr)
    attrs |= {a.name: a.attr for a in attr_array}
  op.attributes["mhlo.frontend_attributes"] = ir.DictAttr.get(attrs)

# TODO(yashkatariya): Delete this after legacy compute_on is deleted.
def wrap_compute_type_in_place(ctx: LoweringRuleContext,
                               op: ir.Value | ir.Operation) -> None:
  if ctx.jaxpr_eqn_ctx is None or ctx.jaxpr_eqn_ctx.compute_type is None:
    return
  op = _get_owner(op)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly one of the supported strings: `device_host`, `device`, or `tpu_sparsecore`
2. Check for typos/dashes vs underscores (e.g. `device-host` is invalid)
3. If you need host vs device vs TPU Sparse Core placement, confirm which of the three supported modes matches your intent in current JAX docs

Example fix

# before
compute_type = 'device-host'

# after
compute_type = 'device_host'

VALID_COMPUTE_TYPES = {'device_host', 'device', 'tpu_sparsecore'}
if compute_type not in VALID_COMPUTE_TYPES:
    raise ValueError(f'compute_type must be one of {VALID_COMPUTE_TYPES}')

def is_valid_compute_type(c: str) -> bool:
    return c in {'device_host', 'device', 'tpu_sparsecore'}

When it happens

Trigger: Passing an invalid string to a JAX API that accepts a compute/compute-type argument (e.g. a Pallas or TPU-related config taking compute type), such as compute_type='gpu', 'host', 'sparse', or a typo like 'device-host'. The value flows down to this normalizer in jax/_src/interpreters/mlir.py.

Common situations: Typos in config strings; assuming older/alternative names like `host` or `sparse` work; copying compute type values from outdated docs or other frameworks.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• Default value must be of type str, got {default} of type {ge
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• Default value must be of type str or None, got {default} of
• new enum value must be None or in {enum_values}, got {new_va

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