[Feature suggestion] Naming einops keras layers

[Sangohe]

Hello, it would be really nice if one could name the eniops layers just like any other keras layer. Right now, the following code triggers an error.

from einops.layers.tensorflow import Rearrange

tf.keras.Sequential([
    tf.keras.Input((224, 224, 3), name="inputs"),
    Rearrange("b h w c -> b c h w", name="rearrange_layer_1"),
])

The error goes away if we do not name the einops layer.

from einops.layers.tensorflow import Rearrange

tf.keras.Sequential([
    tf.keras.Input((224, 224, 3), name="inputs"),
    Rearrange("b h w c -> b c h w"),
])

Naming layers is very useful in keras, specially when using Functional models, to extract intermediate representations or to add new nodes to the graph. This process of extracting nodes is done by accessing the model's layer with model.get_layer(name).

[arogozhnikov]

From looking at implementation I believe you can use

layer._init_set_name('whatever')

Did not check though.