How to use the einops.layers.keras.Rearrange function in einops

def create_model():
            return Sequential([
                Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]),
                Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
                Conv2d(16, kernel_size=5),
                Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
                Rearrange('b c h w -> b (c h w)'),
                Linear(120),
                ReLU(),
                Linear(84),
                ReLU(),
                Linear(10),
            ])