xai-org/x-algorithm · error · ValueError

async_emb axis {axis!r} is not a mesh axis of {mesh}

Error message

async_emb axis {axis!r} is not a mesh axis of {mesh}

What it means

make_context_handle validates that every axis named in table_axis or data_axis exists as an axis of the jax.sharding.Mesh. Sharding embeddings with async_emb requires the mesh to be constructed with those named axes, so an axis name absent from mesh.shape is treated as a programming/config error and rejected before any collective is set up.

Source

Thrown at phoenix/xrex/cuda/async_emb/async_emb.py:103

            emb_width=self.emb_width,
            num_unique=self.num_unique,
            num_devices_per_node=self.num_devices_per_node,
        )


def make_context_handle(
    mesh: jax.sharding.Mesh,
    table_axis: tuple[str, ...],
    *,
    data_axis: tuple[str, ...],
    tokens_per_batch: int,
    emb_width: int,
    num_unique: int,
    num_devices_per_node: int,
) -> AsyncEmbContextHandle:
    for axis in dict.fromkeys((*table_axis, *data_axis)):
        if axis not in mesh.shape:
            raise ValueError(f"async_emb axis {axis!r} is not a mesh axis of {mesh}")
    group_size, flatten_replicas = get_flatten_replica_groups(mesh, table_axis)
    missing_axes = [axis for axis in table_axis if axis not in data_axis]
    if missing_axes:
        raise ValueError(
            f"async_emb requires token shards to vary across the communicator: "
            f"table_axis {missing_axes} missing from data_axis {data_axis}"
        )
    off_communicator_shards = math.prod(
        mesh.shape[axis] for axis in data_axis if axis not in table_axis
    )
    if off_communicator_shards != 1:
        raise ValueError(
            f"async_emb requires exactly one token shard per communicator rank: "
            f"data_axis {data_axis} shards tokens over {off_communicator_shards} "
            f"positions outside table_axis {table_axis}"
        )
    if tokens_per_batch % group_size != 0:
        raise ValueError(

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Construct the JAX Mesh with the axis names used by the config, e.g. Mesh(devices, axis_names=('data','table','replica')).
  2. Or update table_axis/data_axis in the async_emb config to the axis names actually present in mesh.shape.
  3. Print mesh.shape right before training to confirm the available axis names.

Example fix

# before
mesh = jax.sharding.Mesh(jax.devices(), ('batch',))
handle = make_context_handle(mesh, table_axis='table', data_axis='batch', ...)
# ValueError: async_emb axis 'table' is not a mesh axis of ...

# after
devices = jax.devices().reshape(-1, 1)
mesh = jax.sharding.Mesh(devices, axis_names=('batch', 'table'))
handle = make_context_handle(mesh, table_axis='table', data_axis='batch', ...)
Defensive patterns

Strategy: validation

Validate before calling

def check_async_emb_axes(mesh, table_axis, data_axis):
    for axis in dict.fromkeys((*table_axis, *data_axis)):
        if axis not in mesh.shape:
            raise ValueError(
                f"axis {axis!r} missing from mesh {tuple(mesh.shape)}; "
                f"available: {tuple(mesh.axis_names)}"
            )

check_async_emb_axes(mesh, table_axis, data_axis)  # before make_context_handle

Type guard

def axes_are_mesh_axes(mesh, axes) -> bool:
    return all(a in mesh.shape for a in axes)

Try / catch

try:
    handle = make_context_handle(mesh, table_axis, data_axis, ...)
except ValueError as e:
    if "is not a mesh axis" in str(e):
        raise SystemExit(f"fix mesh axis names: {e}") from e
    raise

Prevention

When it happens

Trigger: Calling make_context_handle (or the higher-level _create_async_emb_executables) with table_axis='table' / data_axis='batch' when the Mesh was built with different axis names, or with an anonymous/default mesh whose shape dict lacks those keys.

Common situations: Renaming mesh axes in the launcher without updating the async_emb config; using a mesh constructed for a different model stage; typo'd axis strings; reusing config from another experiment whose mesh layout differs.

Related errors

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/3d115a3e337873d7. Report an issue: GitHub.