jax-ml/jax · error · ValueError

collective_id has to be specified when using a custom barrie

Error message

collective_id has to be specified when using a custom barrier (cannot auto-allocate without lowering context)

What it means

A custom barrier is requested but no collective_id is given and there is no lowering context to auto-allocate one.

Source

Thrown at jax/_src/tpu_custom_call.py:790

        # and the auto-assigned collective ids so far.
        if (collective_id
            in ctx.module_context.pallas_collective_id_mapping.auto.values()):
          raise ValueError(
              f"The manually assigned {collective_id=} in {kernel_name=}"
              " 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":

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Explicitly pass an integer collective_id when using a custom barrier
2. Drop the custom barrier if not needed

Example fix

// before
lower(..., custom_barrier=True)
// after
lower(..., custom_barrier=True, collective_id=100000)

if custom_barrier and collective_id is None:
    raise ValueError('provide collective_id')  # fail fast with your own message

When it happens

Trigger: Calling the lowering path with use_semaphore or custom barrier params but collective_id=None outside a proper lowering context.

Common situations: Using custom barriers in ahead-of-time or manually driven lowering without providing ids.

Understand the failure class

Background: Missing required parameter errors: what 'X is required' and 'the required X param is missing' mean, and how to fix them — this error's family across 27 libraries.

Related errors

• The manually assigned {collective_id=} in {kernel_name=} con
• collective_id has to be unspecified or None when not using a
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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