jax-ml/jax · error · ValueError

explicit opt_level is only supported for SparseCore kernels.

Error message

explicit opt_level is only supported for SparseCore kernels.

What it means

An explicit opt_level was passed for a non-sparsecore kernel; explicit optimization level is only supported for SparseCore.

Source

Thrown at jax/_src/tpu_custom_call.py:809

          "collective_id has to be specified when using a custom barrier "
          "(cannot auto-allocate without lowering context)"
      )
  elif collective_id is not None and not allow_collective_id_without_custom_barrier:
    raise ValueError(
        "collective_id has to be unspecified or None when not using a custom"
        " barrier"
    )
  if vmem_limit_bytes is not None and not isinstance(vmem_limit_bytes, int):
    raise ValueError(
        "vmem_limit_bytes must be an int: provided with a"
        f" {type(vmem_limit_bytes)}."
    )
  if tiling is not None and  device_type != "sparsecore":
    raise ValueError(
        "explicit tiling is only supported for SparseCore kernels."
    )
  if opt_level is not None and device_type != "sparsecore":
    raise ValueError(
        "explicit opt_level is only supported for SparseCore kernels."
    )
  return CustomCallBackendConfig(
      lowered_module_asm,
      lowered_module_asm_version,
      has_communication,
      collective_id,
      device_type,
      cost_estimate,
      needs_hlo_passes,
      needs_layout_passes,
      vmem_limit_bytes,
      flags,
      allow_input_fusion,
      serialization_format,
      internal_scratch_in_bytes,
      output_memory_spaces,
      disable_bounds_checks,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove opt_level for TensorCore kernels
2. Use the default optimization pipeline for dense kernels

Example fix

// before
kernel(..., compiler_params=dict(opt_level=2))
// after
kernel(...)

if device_type != 'sparsecore':
    params.pop('opt_level', None)

When it happens

Trigger: Passing opt_level in compiler params to a dense TensorCore Pallas kernel.

Common situations: Trying to control optimization of dense kernels via opt_level.

Related errors

• explicit tiling is only supported for SparseCore kernels.
• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor

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