mlp_mixer_2.py

# Token Mixing MLPs : Allow communication within patches.
def token_mixing( x , token_mixing_mlp_dims ):
    # x is a tensor of shape ( batch_size , num_patches , channels )
    x = tf.keras.layers.LayerNormalization( epsilon=1e-6 )( x )
    x = tf.keras.layers.Permute( dims=[ 2 , 1 ] )( x ) 
    # After transposition, shape of x -> ( batch_size , channels , num_patches )
    x = mlp( x , token_mixing_mlp_dims )
    return x