flax EinMix syntax differs from torch

EinMix implementation for torch uses **axes_lengths to initialize the axes lengths while the flax implementation requires providing a sizes dictionary. This causes the code required for the two frameworks to differ, but it would be nice if the uniformity einops maintains across all other modules was kept here.

in test_layers.py

torch:

EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84),

jax:

x = EinMix("b (h h2) (w w2) c -> b h w c_out", "h2 w2 c c_out", "c_out", sizes=dict(h2=2, w2=3, c=4, c_out=5))(x)

@arogozhnikov is it an easy fix?