jax-ml/jax · error · ValueError

collective_id has to be unspecified or None when not using a

Error message

collective_id has to be unspecified or None when not using a custom barrier

What it means

collective_id was provided but no custom barrier is used, and allow_collective_id_without_custom_barrier is not enabled.

Source

Thrown at jax/_src/tpu_custom_call.py:795

              " conflicts with an existing auto-assigned collective id."
              " Auto-assignment uses a base collective id of"
              f" {_AUTO_COLLECTIVE_BASE_ID}. Please use values away from this"
              " offset."
          )
        ctx.module_context.pallas_collective_id_mapping.manual[key] = (
            collective_id
        )
        ctx.module_context.pallas_collective_id_mapping.all_ids.add(
            collective_id
        )

    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove collective_id from the params, or enable the custom-barrier option / allow_collective_id_without_custom_barrier flag

Example fix

// before
kernel(..., compiler_params=dict(collective_id=5))
// after
kernel(...)  # no collective_id without a barrier

assert collective_id is None or using_custom_barrier or allow_flag

When it happens

Trigger: Passing collective_id in kernel compiler params without enabling the corresponding barrier option.

Common situations: Copy-pasting barrier config into kernels that don't use barriers.

Related errors

• The manually assigned {collective_id=} in {kernel_name=} con
• collective_id has to be specified when using a custom barrie
• explicit tiling is only supported for SparseCore kernels.
• explicit opt_level is only supported for SparseCore kernels.
• megacore_mode must be one of ['kv_head', 'batch', None]

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