feat: add Repeat layer to einops.layers across backends - #439
Open
mvanhorn wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a
Repeatlayer toeinops.layers, complementing the existingRearrangeandReducelayers.Repeatwraps the functionaleinops.repeatso it can be dropped intonn.Sequential-style model definitions. It is implemented across every backend (torch, tensorflow, keras, flax, oneflow, paddle) exactly parallel to howReduceis defined, and follows theReducetemplate: aRepeatMixin(ReduceMixin)ineinops/layers/__init__.pywith the reduction hard-wired to"repeat", plus a concreteRepeatclass in each backend layer module.Because
Repeatneeds noreductionargument, the mixin overrides__repr__,__getstate__, and__setstate__to drop that field, so it pickles and round-trips likeRearrange.Why this matters
Users have long asked for a first-class
Repeatlayer parallel toRearrangeandReduce(#185). Today the only layer-based option isReduce('a b -> a b c', reduction='repeat', c=4), which reads confusingly because a repeat is not a reduction. The maintainer green-lit incorporating aRepeatlayer in the issue thread, and multiple users noted concrete use cases, includingtorch.jit.scriptcompatibility that only the layer forms provide. This change mirrors the acceptedReducetemplate across all backends and keeps the torch pathtorch.jit.script-safe (for a repeat,reduced_axesis empty and only theadd_axesbranch runs).Testing
Added layer tests parallel to the existing
Reducetests, covering numpy and torch: imperative and symbolic parity against the functionaleinops.repeat, pickle round-trips, wrong-shape rejection, an invalid-pattern error path, and atorch.jit.scripttest that scripts aRepeatinside aSequentialand checks scripted output matches eager. The other backends are exercised by the same parametrized layer tests in CI.Fixes #185