a0f77cc8b87512e691a7c22b6fb9ba60f603b632,layer/reparameterization_trick.py,ReparameterizationLayer,layer_op,#ReparameterizationLayer#Any#,36

Before Change

            stochastic_parts = []
            for p in range(0, self.number_of_samples):
                stochastic_parts.append(noise_like(means, self.prior))
            noise_sample = tf.reduce_mean(tf.stack(stochastic_parts), axis=0)

        output_tensor = means + tf.exp(0.5 * logvariances) * noise_sample

After Change

        if self.number_of_samples == 1:
            noise_sample = noise_shaped_like(tf.shape(means), self.prior)
        else:
            shape_of_expanded_sample = tf.concat([tf.constant(self.number_of_samples, shape=[1, ]), tf.shape(means)], 0)
            noise_sample = noise_shaped_like(shape_of_expanded_sample, self.prior)
            noise_sample = tf.reduce_mean(noise_sample, axis=0)

        return means + tf.exp(0.5 * logvariances) * noise_sample

Instances