import tensorflow as tf

            # Perform affine mapping at each layer of the neural network
            Z = tf.layers.dense(Z, n_basis//2)
            # Define variational parameters
            alpha_mean = tf.get_variable('alpha_mean_layer'+str(h),
                shape=[1, 1, n_basis, n_out],
                initializer=tf.random_normal_initializer())
            alpha_logstd = tf.get_variable('alpha_logstd_layer'+str(h),
                shape=[1, 1, n_basis, n_out],
                initializer=tf.random_normal_initializer())
            alpha_std = tf.exp(alpha_logstd)
