import tensorflow as tf

        elif type=='mix_scramble':
            self.augment = self.mix_scramble  
        elif type=='blur':
            self.augment = self.gaussian_blur
            self.pointwise_filter = tf.eye(3, batch_shape=[1, 1])

        elif type=='high_low_pass':
            self.augment = self.high_low_pass
            self.kernel = self.gaussian_kernel(size,mean,std)
            self.kernel = tf.tile(self.kernel[:, :, tf.newaxis, tf.newaxis], [1, 1, 3, 1])
            self.pointwise_filter = tf.eye(3, batch_shape=[1, 1])
            self.paddings = [[size,size],[size,size],[0,0]]
        elif type=='no_op':
            self.augment = self.no_op


    def gaussian_kernel(self,size,mean,std):
        """Makes 2D gaussian Kernel for convolution."""
        d = tfp.distributions.Normal(mean, std)
        vals = d.prob(tf.range(start = -size, limit = size + 1, dtype = tf.float32))
