radi-cho

# x_spatial shape: [batch_size, num_patches, embedding_dim].
x_spatial = self.spatial_gating_unit(x_projected)
# x_projected shape: [batch_size, num_patches, embedding_dim].
x_projected = self.channel_projection2(x_spatial)

radi-cho

def call(self, inputs):
    x = self.normalize1(inputs)
    x_projected = self.channel_projection1(x)

radi-cho

v_channels = tf.linalg.matrix_transpose(v)
v_projected = self.spatial_projection(v_channels)
v_projected = tf.linalg.matrix_transpose(v_projected)

radi-cho

self.spatial_projection = layers.Dense(units=num_patches, bias_initializer="Ones")

radi-cho

def spatial_gating_unit(self, x):
    # u and v shape: [batch_size, num_patchs, embedding_dim]
    u, v = tf.split(x, num_or_size_splits=2, axis=2)
    v = self.normalize2(v)

radi-cho

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class gMLPLayer(layers.Layer):
    def __init__(self, num_patches, embedding_dim, dropout_rate, *args, **kwargs):
        super(gMLPLayer, self).__init__(*args, **kwargs)

radi-cho

self.normalize1 = layers.LayerNormalization(epsilon=1e-6)
self.normalize2 = layers.LayerNormalization(epsilon=1e-6)

self.channel_projection1 = keras.Sequential(
    [
        layers.Dense(units=embedding_dim * 2),
        layers.ReLU(),
        layers.Dropout(rate=dropout_rate),
    ]
)

radi-cho

tf.one_hot(tf.argmax(p), depth = len(p))

<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0., 1., 0.], dtype=float32)>

radi-cho

temperature = 0.01
dist = tfp.distributions.RelaxedOneHotCategorical(temperature, probs=p)
dist.sample()

<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0., 1., 0.], dtype=float32)>

radi-cho

temperature = 10
dist = tfp.distributions.RelaxedOneHotCategorical(temperature, probs=p)
dist.sample()

<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.31916314, 0.34642866, 0.33440822], dtype=float32)>