{"record":{"id":"3d115a3e337873d7","repo":"xai-org/x-algorithm","slug":"async-emb-axis-axis-r-is-not-a-mesh-axis-of-mes","errorCode":null,"errorMessage":"async_emb axis {axis!r} is not a mesh axis of {mesh}","messagePattern":"async_emb axis (.+?) is not a mesh axis of (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/async_emb/async_emb.py","lineNumber":103,"sourceCode":"            emb_width=self.emb_width,\n            num_unique=self.num_unique,\n            num_devices_per_node=self.num_devices_per_node,\n        )\n\n\ndef make_context_handle(\n    mesh: jax.sharding.Mesh,\n    table_axis: tuple[str, ...],\n    *,\n    data_axis: tuple[str, ...],\n    tokens_per_batch: int,\n    emb_width: int,\n    num_unique: int,\n    num_devices_per_node: int,\n) -> AsyncEmbContextHandle:\n    for axis in dict.fromkeys((*table_axis, *data_axis)):\n        if axis not in mesh.shape:\n            raise ValueError(f\"async_emb axis {axis!r} is not a mesh axis of {mesh}\")\n    group_size, flatten_replicas = get_flatten_replica_groups(mesh, table_axis)\n    missing_axes = [axis for axis in table_axis if axis not in data_axis]\n    if missing_axes:\n        raise ValueError(\n            f\"async_emb requires token shards to vary across the communicator: \"\n            f\"table_axis {missing_axes} missing from data_axis {data_axis}\"\n        )\n    off_communicator_shards = math.prod(\n        mesh.shape[axis] for axis in data_axis if axis not in table_axis\n    )\n    if off_communicator_shards != 1:\n        raise ValueError(\n            f\"async_emb requires exactly one token shard per communicator rank: \"\n            f\"data_axis {data_axis} shards tokens over {off_communicator_shards} \"\n            f\"positions outside table_axis {table_axis}\"\n        )\n    if tokens_per_batch % group_size != 0:\n        raise ValueError(","sourceCodeStart":85,"sourceCodeEnd":121,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/async_emb/async_emb.py#L85-L121","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Construct the JAX Mesh with the axis names used by the config, e.g. Mesh(devices, axis_names=('data','table','replica')).","Or update table_axis/data_axis in the async_emb config to the axis names actually present in mesh.shape.","Print mesh.shape right before training to confirm the available axis names."],"exampleFix":"# before\nmesh = jax.sharding.Mesh(jax.devices(), ('batch',))\nhandle = make_context_handle(mesh, table_axis='table', data_axis='batch', ...)\n# ValueError: async_emb axis 'table' is not a mesh axis of ...\n\n# after\ndevices = jax.devices().reshape(-1, 1)\nmesh = jax.sharding.Mesh(devices, axis_names=('batch', 'table'))\nhandle = make_context_handle(mesh, table_axis='table', data_axis='batch', ...)","handlingStrategy":"validation","validationCode":"def check_async_emb_axes(mesh, table_axis, data_axis):\n    for axis in dict.fromkeys((*table_axis, *data_axis)):\n        if axis not in mesh.shape:\n            raise ValueError(\n                f\"axis {axis!r} missing from mesh {tuple(mesh.shape)}; \"\n                f\"available: {tuple(mesh.axis_names)}\"\n            )\n\ncheck_async_emb_axes(mesh, table_axis, data_axis)  # before make_context_handle","typeGuard":"def axes_are_mesh_axes(mesh, axes) -> bool:\n    return all(a in mesh.shape for a in axes)","tryCatchPattern":"try:\n    handle = make_context_handle(mesh, table_axis, data_axis, ...)\nexcept ValueError as e:\n    if \"is not a mesh axis\" in str(e):\n        raise SystemExit(f\"fix mesh axis names: {e}\") from e\n    raise","preventionTips":["Define axis names once in a shared layout module and derive both the Mesh and async_emb config from it.","Log mesh.shape at job start to catch renames before they hit the collective setup."],"tags":["jax","sharding","mesh-axes","distributed","async-embedding"],"backgroundTag":"jax-mesh-axis-validation","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}