jax-ml/jax · error · ValueError

Only one of `core_axis` or `core_axis_name` can be set.

Error message

Only one of `core_axis` or `core_axis_name` can be set.

What it means

Multi-core grid partitioning can be specified either positionally via core_axis or by name via core_axis_name. Supplying both is ambiguous, so the API rejects it with ValueError.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2023

    out_specs=(),
    tiling: Tiling | None = None,
    core_axis: tuple[int, ...] | int | None = None,
    core_axis_name: tuple[str, ...] | str | None = None,
    dimension_semantics: tuple[GridDimensionSemantics, ...] | None = None,
    trace_scopes: bool = True,
    no_pipelining: bool = False,
    _explicit_indices: bool = False,
):
  in_specs = _normalize_specs(in_specs)
  out_specs = _normalize_specs(out_specs)

  if any(g <= 0 for g in grid if isinstance(g, int)):
    raise ValueError(
        f"All elements in the grid must be strictly positive, but got {grid=}"
    )

  if core_axis is not None and core_axis_name is not None:
    raise ValueError("Only one of `core_axis` or `core_axis_name` can be set.")
  core_axis_ = core_axis_name if core_axis is None else core_axis
  if dimension_semantics is None:
    dimension_semantics = (ARBITRARY,) * len(grid)

  num_in_specs = len(in_specs)

  def wrapped(*args, allocations=None, **kwargs):
    num_cores, core_id = _resolve_core_info(core_axis_)

    if allocations is not None and not in_specs and not out_specs:
      flat_allocs = [
          b for b in allocations if isinstance(b, BufferedRefBase) or b is None
      ]
      in_specs_ = tuple(
          b.spec if isinstance(b, BufferedRefBase) else None
          for b in flat_allocs if b is None or b.buffer_type == BufferType.INPUT
      )
      out_specs_ = tuple(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep exactly one of core_axis or core_axis_name; prefer core_axis_name for clarity
2. If defaults leaked in from a config dict, unset the other key before passing kwargs

Example fix

# before
emit_pipeline(..., core_axis=1, core_axis_name='data')
# after
emit_pipeline(..., core_axis_name='data')

assert not (core_axis is not None and core_axis_name is not None), \
    'pass only one of core_axis / core_axis_name'

When it happens

Trigger: Calling emit_pipeline with both core_axis=1 and core_axis_name='data' set on the same invocation.

Common situations: Refactoring from positional to named axis specification and leaving both arguments populated; copy-pasted kwargs from another call site.

Related errors

• Sum of sizes {n} must be equal to dimension {axis} of the op
• Cannot pull iota along dimension {dimension} with None block
• cannot specify both devices and num_cores
• Invalid memory space: {memory_space!r}
• dimension_semantics must be the same length as grid.

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