jax-ml/jax · error · ValueError

explicit tiling is only supported for SparseCore kernels.

Error message

explicit tiling is only supported for SparseCore kernels.

What it means

An explicit tiling parameter was passed for a kernel whose device_type is not sparsecore; explicit tiling only applies to SparseCore kernels.

Source

Thrown at jax/_src/tpu_custom_call.py:805

        )

    if collective_id is None:
      raise ValueError(
          "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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the tiling param for non-sparsecore kernels
2. Or ensure the kernel is actually a SparseCore kernel (target sparsecore)

Example fix

// before
kernel(..., compiler_params=dict(tiling=tiling))
// after
kernel(...)  # omit tiling for TensorCore kernels

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

When it happens

Trigger: Passing tiling=... in compiler params to a TensorCore (dense) Pallas kernel.

Common situations: Reusing sparse-core config on dense kernels.

Related errors

• explicit opt_level 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/fc374c8d6cbe1eca. Report an issue: GitHub.