Alex Rogozhnikov

@SauravMaheshkar here is an example for torch: https://github.com/arogozhnikov/einops/blob/master/tests/test_layers.py#L200-L231 Goal of adding tests is to make sure that layers can be serialized / deserialized using native flax serialization mechanisms

hm. Yes, that's tricky, flax doesn't have own backend. ``` try: import flax has_flax = True except: has_flax = False ```

Hmmm, `torch.fx` is ... awkward. Using symboling tracing in a framework that always bet on define-by-run is breaking rules. Einops requires some information about shape, but torch.fx does not provide...

Hi @lkhphuc Situation is: - I'd like to have jax layers, but reasonably I can support one or maybe two layer systems, no more. - there are a dozen (maybe...

@SauravMaheshkar yes, sure

Hi @JBOE22175, thanks for very detailed description of issue and the notebook. - layers that take pattern as a part of forward - this doesn't look like "designed" usage for...

see #172 , that's a plan for the next release

Hi Christoph, thank you for following issue template, despite it did not fit. Please read a relevant discussion here: https://github.com/arogozhnikov/einops/discussions/126 As far as I see, in most cases repeat can...

> you could also say, that you never need the function repeat not really, np.repeat and np.tile behavior can't be obtained from other functions. It's just my observation that use-cases...

Thank you, always glad to hear that einops is helpful. > I hacked up a prototype here https://github.com/lucidrains/memorizing-transformers-pytorch/commit/b190c5ec6144d6dafb11eb356fe51dd10bccc052 to give you an idea of the use-case. What you suggest with...