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.
zeta/ops/expand.pyfails on its own__main__example:There are two problems:
1.
combined_dimsproduces garbage axis names. The function zips the raw whitespace-split input pattern with the tensor shape (line 37):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 vianew_dims.2. The example's sizes are inconsistent.
pattern = "b (c h) (w f) -> b c h w"withc=4, h=25implies the second axis isc*h = 100, but the tensor's second dim is50. Alsof=8is passed butfdoesn't appear in the output, so it's flagged as one-sided.A minimal version that works:
Happy to send a PR if you'd like.