jax-ml/jax · error · ValueError

The manually assigned {collective_id=} in {kernel_name=} con

Error message

The manually assigned {collective_id=} in {kernel_name=} conflicts with an existing auto-assigned collective id. Auto-assignment uses a base collective id of {_AUTO_COLLECTIVE_BASE_ID}. Please use values away from this offset.

What it means

A user-specified collective_id collides with an id the compiler auto-assigned (auto ids start at _AUTO_COLLECTIVE_BASE_ID) for barriers in the same module.

Source

Thrown at jax/_src/tpu_custom_call.py:775

              _AUTO_COLLECTIVE_BASE_ID + auto_num + len(existing_ids) + 1)
          new_id = next(id for id in proposed_ids if id not in existing_ids)
          ctx.module_context.pallas_collective_id_mapping.auto[key] = new_id
          ctx.module_context.pallas_collective_id_mapping.all_ids.add(new_id)
          collective_id = new_id
          if (len(ctx.module_context.pallas_collective_id_mapping.auto)
              > _AUTO_COLLECTIVE_ID_LIMIT):
            logging.warning(
                "The number of auto-assigned collective ids for pallas kernels"
                " is very large, consider manually annotating the kernels with"
                " collective ids:"
                f" {ctx.module_context.pallas_collective_id_mapping}"
            )
      else:  # Manually assigned collective ID.
        # We need to check for a conflict between the manual collective id
        # 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)"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pick a collective_id far from _AUTO_COLLECTIVE_BASE_ID (e.g. a large number like 10000+)
2. Avoid specifying collective_id and let the compiler auto-assign

Example fix

// before
kernel(..., compiler_params=dict(collective_id=50))
// after
kernel(..., compiler_params=dict(collective_id=100000))

from jax._src.tpu_custom_call import _AUTO_COLLECTIVE_BASE_ID as BASE
assert collective_id is None or collective_id >= BASE + 10000

When it happens

Trigger: Passing a manual collective_id to a Pallas kernel that falls in the auto-assigned range while other collectives were auto-allocated in the same lowering.

Common situations: Choosing small collective_id values that overlap the auto base offset when multiple kernels use custom barriers.

Related errors

• collective_id has to be specified when using a custom barrie
• 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/ae4b9b2a00c97354. Report an issue: GitHub.