TensorFlow 2.0 implementation of a sampling layer for a variational autoencoder.

sampling.py

class Sampling(tf.keras.layers.Layer):
  def call(self, args):
    z_mean, z_log_var = args
    batch = tf.shape(z_mean)[0]
    dim = tf.shape(z_mean)[1]
    epsilon = tf.random.normal(shape=(batch, dim), mean=0., stddev=1.)
    return z_mean + epsilon * tf.exp(0.5 * z_log_var)