Add dict syntax to the einops layers example

[eadadi]

This patch adds an example of the dict syntax that is given as argument in einops layers.

Why do we need this change?
I was trying to make an einops layer doing the following rearrange:

Rearrange('b (h w c) -> b h w c', h=64, w=64)

And had to read the source code to discover I have to use dict syntax to make it works:

Rearrange('b (h w c) -> b h w c', {'h'=64, 'w'=64})

So I believe that adding this example to the docs will make it easier for people like me

[arogozhnikov]

@eadadi

Rearrange('b (h w c) -> b h w c', h=64, w=64)  # works for almost all frameworks
Rearrange('b (h w c) -> b h w c', {'h': 64, 'w': 64}) # specific syntax for flax, because flax doesn't accept **kwargs

so no, this change is incorrect.

[eadadi]

ok good to know. modified accordingly