{"record":{"id":"a37c7fc9b911b1a7","repo":"xai-org/x-algorithm","slug":"async-emb-requires-token-shards-to-vary-across-the","errorCode":null,"errorMessage":"async_emb requires token shards to vary across the communicator: table_axis {missing_axes} missing from data_axis {data_axis}","messagePattern":"async_emb requires token shards to vary across the communicator: table_axis (.+?) missing from data_axis (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/async_emb/async_emb.py","lineNumber":107,"sourceCode":"\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(\n            f\"async_emb tokens_per_batch={tokens_per_batch} does not shard evenly \"\n            f\"over the {group_size}-rank communicator\"\n        )\n    if emb_width % group_size != 0:","sourceCodeStart":89,"sourceCodeEnd":125,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/async_emb/async_emb.py#L89-L125","documentation":"make_context_handle requires every table_axis (embedding-sharding axis defining the communicator group) to also appear in data_axis (token-sharding axes). If tokens are not sharded along an axis on which tables are sharded, ranks in one communicator would hold embeddings for tokens they never see, which the async_emb collective protocol cannot handle.","triggerScenarios":"Passing table_axis=('table','model') but data_axis=('batch',) — the 'table'/'model' axes are missing from data_axis — when calling make_context_handle / _create_async_emb_executables.","commonSituations":"Configuring data parallelism only ('batch') while enabling N-way model-parallel embedding sharding; reusing a single-tower mesh config for the two-tower async_emb path; copy-paste of axis tuples between experiments with different sharding layouts.","solutions":["Add every table_axis name into data_axis, e.g. table_axis=('table',), data_axis=('batch','table'), so token shards vary across the communicator.","If you did not intend to shard tables, reduce table_axis to a single axis that is already in data_axis.","Cross-check the axis names against mesh.shape (see also the adjacent mesh-axis validation errors)."],"exampleFix":"# before\nhandle = make_context_handle(\n    mesh, table_axis=('table',), data_axis=('batch',), ...\n)\n\n# after\nhandle = make_context_handle(\n    mesh, table_axis=('table',), data_axis=('batch', 'table'), ...\n)","handlingStrategy":"validation","validationCode":"def check_axis_containment(table_axis, data_axis):\n    missing = [a for a in table_axis if a not in data_axis]\n    if missing:\n        raise ValueError(\n            f\"table_axis {missing} must also appear in data_axis {data_axis} \"\n            f\"so token shards vary across the communicator\"\n        )\n\ncheck_axis_containment(table_axis, data_axis)  # before make_context_handle","typeGuard":"def table_axes_in_data(table_axis, data_axis) -> bool:\n    return set(table_axis) <= set(data_axis)","tryCatchPattern":"try:\n    handle = make_context_handle(mesh, table_axis, data_axis, ...)\nexcept ValueError as e:\n    if \"missing from data_axis\" in str(e):\n        data_axis = tuple(dict.fromkeys((*data_axis, *table_axis)))\n        handle = make_context_handle(mesh, table_axis, data_axis, ...)\n    else:\n        raise","preventionTips":["Standardize on data_axis = table_axis + one batch axis for async_emb runs.","Add a config-lint step that reifies axis invariants before launching multi-host jobs."],"tags":["jax","sharding","mesh-axes","distributed","async-embedding"],"backgroundTag":"jax-sharding-axis-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}