xai-org/x-algorithm · error · ValueError

Unknown named dimension

Error message

Unknown named dimension: {named_dim} for namespace: {namespace}

What it means

When building a ShardingSpec, _sharding_rule looks up the requested named dimension inside the per-namespace config dict. If the dimension name (e.g. 'in_dim') is absent from that namespace's config, it raises this ValueError naming both the dimension and the namespace.

Solutions

1. Add the missing named dimension with its PartitionSpec to the namespace's config in the sharding config
2. Verify spelling/casing of the dimension against what the model code requests
3. If the dim should not be sharded, add an explicit entry mapping it to replicated (None)

Example fix

# before
sharding:
  attention:
    in_dim: ["data"]
# after
sharding:
  attention:
    in_dim: ["data"]
    out_dim: ["model"]

for ns, dims in cfg['sharding'].items():
    assert required_dims[ns] <= set(dims), f'{ns} missing {required_dims[ns] - set(dims)}'

try:
    spec = rule(named_dim, config, namespace)
except ValueError as e:
    if 'Unknown named dimension' in str(e):
        return None  # replicate
    raise

When it happens

Trigger: A layer asks for a named dim like 'out_dim' under namespace 'attention' but the sharding config for 'attention' only defines 'in_dim'; renaming a dimension in model code without updating the sharding config.

Common situations: Model refactors that rename tensor dimensions; configs written for a different model variant missing entries for some namespaces; typos in dimension names.

Related errors

• is already exist in ShardingContext
• async_emb axis is not a mesh axis of
• async_emb emb_width=
• async_emb requires exactly one token shard per communicator…
• async_emb requires token shards to vary across the…

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

Appendix: source

Thrown at phoenix/xrex/models/sharding_context.py:132

            ]
        flattened_physical_axes, _ = jax.tree.flatten(physical_axes)
        return axis_group_size(flattened_physical_axes, self.mesh)


PerNamespaceShardingConfig = dict[str, ShardingSpec]
ShardingConfig = dict[str, PerNamespaceShardingConfig]


def make_sharding_context_from_config(
    name: str, mesh: Mesh, config: ShardingConfig
) -> ShardingContext:
    ctx = ShardingContext(name, mesh)

    def _sharding_rule(
        named_dim: str, config: PerNamespaceShardingConfig = None, namespace: str = None
    ) -> ShardingSpec:
        if named_dim not in config:
            raise ValueError(f"Unknown named dimension: {named_dim} for namespace: {namespace}")
        return config[named_dim]

    for namespace, c in config.items():
        ctx.register_sharding_rule(namespace)(
            partial(_sharding_rule, config=c, namespace=namespace)
        )
    return ctx


default_sharding_config = {
    "default": {
        "batch": ("expert", "replica", "data"),
        "batch_attn": ("expert", "replica", "data"),
        "dense_activation_model": "model",
        "dense_activation_seq": "seq",
        "embed": None,
        "head": ("seq", "model"),
        "hidden": None,

View on GitHub (pinned to 24c60942c5)