Random Categorical Selection with Gradient capability

random_categorical_layer.py

class RandomCategorical(L.Layer):
  def __init__(self, num_classes=10, factor=0.999):
    super(RandomCategorical, self).__init__()
    self.num_classes = num_classes
    self.factor = factor

  @tf.function 
  def call(self, input):
    sample = tf.random.categorical(logits=input, num_samples=1)
    hot = tf.reshape( tf.one_hot( sample, depth=self.num_classes ), shape=[-1,self.num_classes] )
    return hot * self.factor  + tf.nn.sigmoid(input) * (1. - self.factor)