text: import tensorflow as tf
  
  
  class GaussianSampling(tf.keras.layers.Layer):
      """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""
  
      def __init__(self, name: str = "GaussianSampling"):
          super(GaussianSampling, self).__init__(name=name)
  
      @tf.function
      def call(self, inputs, training: bool = True, test_sample_mode='mean'):
          if training or test_sample_mode=='posterior_sample':
              batch = tf.shape(inputs)[-3]
              dim = tf.shape(inputs)[-2]
              epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
              return inputs[:, :, 0] + tf.exp(0.5 * inputs[:, :, 1]) * epsilon
          else:
              return inputs[:, :, 0]
  