feat: Add Flax Layers

[SauravMaheshkar]

Adds Flax Module's for einops operations. More details can be found in #153

einops supports JAX arrays, so I'm unsure if I should add tests or not ?

CC: @arogozhnikov

[arogozhnikov]

I should add tests or not

Yes, tests are needed, because we need to ensure that following things are working:

1. serialization/deserialization of weights
2. serialization/deserialization/pickling/deepcopy-ing of models
3. layers normally interact with other layers

For other frameworks, I just have some lenet-style model with layers, and check that 1 and 2 work for the whole model

[arogozhnikov]

@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

[SauravMaheshkar]

Request for Clarification:

has_flax = any(backend.framework_name == 'jax' for backend in collect_test_backends(symbolic=False, layers=True))

would be the correct way ?

[arogozhnikov]

hm. Yes, that's tricky, flax doesn't have own backend.

try:
    import flax
    has_flax = True
except:
    has_flax = False

[arogozhnikov]

Hi @SauravMaheshkar I think something was changed in the flax, because the code does not work.

I have another implementation #214 that propbably accounts for all strangeness of flax, so closing this. Testing is very welcome!