Skip to content

feat: add Repeat layer to einops.layers across backends - #439

Open
mvanhorn wants to merge 1 commit into
arogozhnikov:mainfrom
mvanhorn:feat/185-repeat-layer
Open

feat: add Repeat layer to einops.layers across backends#439
mvanhorn wants to merge 1 commit into
arogozhnikov:mainfrom
mvanhorn:feat/185-repeat-layer

Conversation

@mvanhorn

@mvanhorn mvanhorn commented Jul 5, 2026

Copy link
Copy Markdown

Summary

Adds a Repeat layer to einops.layers, complementing the existing Rearrange and Reduce layers. Repeat wraps the functional einops.repeat so it can be dropped into nn.Sequential-style model definitions. It is implemented across every backend (torch, tensorflow, keras, flax, oneflow, paddle) exactly parallel to how Reduce is defined, and follows the Reduce template: a RepeatMixin(ReduceMixin) in einops/layers/__init__.py with the reduction hard-wired to "repeat", plus a concrete Repeat class in each backend layer module.

Because Repeat needs no reduction argument, the mixin overrides __repr__, __getstate__, and __setstate__ to drop that field, so it pickles and round-trips like Rearrange.

Why this matters

Users have long asked for a first-class Repeat layer parallel to Rearrange and Reduce (#185). Today the only layer-based option is Reduce('a b -> a b c', reduction='repeat', c=4), which reads confusingly because a repeat is not a reduction. The maintainer green-lit incorporating a Repeat layer in the issue thread, and multiple users noted concrete use cases, including torch.jit.script compatibility that only the layer forms provide. This change mirrors the accepted Reduce template across all backends and keeps the torch path torch.jit.script-safe (for a repeat, reduced_axes is empty and only the add_axes branch runs).

Testing

Added layer tests parallel to the existing Reduce tests, covering numpy and torch: imperative and symbolic parity against the functional einops.repeat, pickle round-trips, wrong-shape rejection, an invalid-pattern error path, and a torch.jit.script test that scripts a Repeat inside a Sequential and checks scripted output matches eager. The other backends are exercised by the same parametrized layer tests in CI.

EINOPS_TEST_BACKENDS=numpy,torch python -m pytest einops/tests/test_layers.py -q
13 passed, 2 skipped

Fixes #185

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Missing einops.layers.....Repeat

1 participant