jax-ml/jax · error · ValueError

cannot specify both devices and num_cores

Error message

cannot specify both devices and num_cores

What it means

create_tensorcore_mesh builds a TensorCoreMesh from either an explicit device sequence or a core count. Specifying both devices and num_cores is ambiguous and rejected outright.

Source

Thrown at jax/_src/pallas/mosaic/core.py:450

    return [
        MemorySpace.VMEM,
        MemorySpace.SMEM,
        MemorySpace.CMEM,
        MemorySpace.SEMAPHORE,
    ]

  @contextlib.contextmanager
  def tracing_context(self):
    yield


def create_tensorcore_mesh(
    axis_name: str,
    devices: Sequence[jax.Device] | None = None,
    num_cores: int | None = None,
) -> TensorCoreMesh:
  if devices is not None and num_cores is not None:
    raise ValueError("cannot specify both devices and num_cores")
  if num_cores is None:
    if devices is None:
      num_cores = _get_default_num_cores()
    else:
      num_cores = devices[0].num_cores
  return TensorCoreMesh(axis_name=axis_name, num_cores=num_cores)


def _convert_semaphore_type_to_aval(
    out_shape: SemaphoreType,
) -> jax_core.AbstractValue:
  return out_shape.get_array_aval()


pallas_core._out_shape_to_aval_mapping[SemaphoreType] = (
    _convert_semaphore_type_to_aval
)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass exactly one of devices or num_cores
2. If using devices, omit num_cores — it is derived from devices[0].num_cores
3. Audit wrapper functions that forward **kwargs to this API

Example fix

// before
mesh = create_tensorcore_mesh('tc', devices=devs, num_cores=4)
// after
mesh = create_tensorcore_mesh('tc', devices=devs)

assert (devices is None) != (num_cores is None), 'pass exactly one of devices/num_cores'

When it happens

Trigger: Calling jax._src.pallas.mosaic.core.create_tensorcore_mesh(axis_name, devices=[d], num_cores=4) — both keyword arguments supplied.

Common situations: Copy-pasted configuration evolving from num_cores to devices-based setup leaving both args in place; wrapper functions forwarding **kwargs that include both.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• You can't use two different TensorCoreMeshes.
• When passing host local inputs to pjit, devices connected to
• Failed to find assignment for logical_axis_index {logical_ax
• Compiler params for platform {platform} cannot be used for {

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