Skip to content

expand() builds invalid einops axis names from input_pattern.split(); example crashes #407

Description

@thehalleyyoung

zeta/ops/expand.py fails on its own __main__ example:

einops.EinopsError:  Error while processing rearrange-reduction pattern "b (c h) (w f) -> b c h w".
 Input tensor shape: torch.Size([2, 50, 64]). Additional info: {'c': 4, 'h': 25, 'w': 8, 'f': 8, 'b': 2, '(c': 50, 'h)': 64}.
 Identifiers only on one side of expression (should be on both): {'f'}

There are two problems:

1. combined_dims produces garbage axis names. The function zips the raw whitespace-split input pattern with the tensor shape (line 37):

combined_dims = {
    **new_dims,
    **dict(zip(input_pattern.split(), tensor.shape)),
}

For "b (c h) (w f)", input_pattern.split() is ['b', '(c', 'h)', '(w', 'f)'], so the dict gets keys like '(c' and 'h)' (visible in the error's "Additional info"). Those aren't valid einops axis names. einops already infers the sizes of un-parenthesized axes from the tensor shape, so this zip line is both wrong and unnecessary — for grouped axes you only need to supply the split sizes via new_dims.

2. The example's sizes are inconsistent. pattern = "b (c h) (w f) -> b c h w" with c=4, h=25 implies the second axis is c*h = 100, but the tensor's second dim is 50. Also f=8 is passed but f doesn't appear in the output, so it's flagged as one-sided.

A minimal version that works:

def expand(tensor, pattern, **new_dims):
    input_pattern, output_pattern = (p.strip() for p in pattern.split("->"))
    return rearrange(tensor, f"{input_pattern} -> {output_pattern}", **new_dims)

# tensor (2, 50, 64); 50 = 2*25, 64 = 8*8
expand(torch.randn(2, 50, 64), "b (c h) (w f) -> b c h w", c=2, h=25, w=8, f=8)

Happy to send a PR if you'd like.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions