jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _normalize_labels(self, labels):
        """Normalize labels"""
        labels = tf.expand_dims(tf.cast(labels, tf.float32), -1)
        labels = tf.divide(labels, 255)

        labels = tf.image.resize_images(labels, (388, 388))
        labels = tf.image.resize_image_with_crop_or_pad(labels, 572, 572)

        cond = tf.less(labels, 0.5 * tf.ones(tf.shape(labels)))
        labels = tf.where(cond, tf.zeros(tf.shape(labels)), tf.ones(tf.shape(labels)))

        return tf.one_hot(tf.squeeze(tf.cast(labels, tf.int32)), 2)
//  [python]
def __call__(self,x,y,stable=None):
        stable = self.stable if stable is None else stable
        if stable:
            x = not_zeros(x)
        return tf.where(tf.less_equal(x,y),tf.ones_like(x),tf.truediv(y,x))
//  [python]
def __call__(self,x,y,stable=None):
        stable = self.stable if stable is None else stable
        if stable:
            x = not_zeros(x)
        return tf.where(tf.less_equal(x,y),tf.ones_like(x),tf.divide(y,x))
//  [python]
def compute_barycentric_coordinates(points, fixed_mesh_vertices, fixed_mesh_faces):
    n_points = points.shape[0]
    n_faces = fixed_mesh_faces.shape[0]
    fixed_mesh_vertices = tf.concat([fixed_mesh_vertices, tf.zeros([n_points, 1])], axis = 1)
    points = tf.concat([points, tf.zeros([n_points, 1])], axis = 1)
    p_a = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 0])[:, tf.newaxis,:], [1, n_points, 1])
    p_b = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 1])[:, tf.newaxis,:], [1, n_points, 1])
    p_c = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 2])[:, tf.newaxis,:], [1, n_points, 1])
    points = tf.tile(points[tf.newaxis,: ,:], [n_faces, 1, 1])

    area_ABC = tf.abs(tf.linalg.cross(p_b -p_a, p_c - p_a)[:, :, 2])
    area_PBC = tf.abs(tf.linalg.cross(p_b -points, p_c - points)[:, :, 2])
    area_PAB = tf.abs(tf.linalg.cross(p_b -points, p_a - points)[:, :, 2])
    area_PCA = tf.abs(tf.linalg.cross(p_c -points, p_a - points)[:, :, 2])

    u= tf.truediv(area_PBC, area_ABC)
    v= tf.truediv(area_PCA, area_ABC)
    w= tf.truediv(area_PAB, area_ABC)
    mask = u + w + v>1+ 1e-6
    u = tf.where(mask, tf.zeros_like(u), u)
    v = tf.where(mask, tf.zeros_like(v), v)
    w = tf.where(mask, tf.zeros_like(w), w)

    return v, w, u
//  [python]
def compute_barycentric_coordinates(points, fixed_mesh_vertices, fixed_mesh_faces):
    n_points = points.shape[0]
    n_faces = fixed_mesh_faces.shape[0]
    fixed_mesh_vertices = tf.concat([fixed_mesh_vertices, tf.zeros([n_points, 1])], axis = 1)
    points = tf.concat([points, tf.zeros([n_points, 1])], axis = 1)
    p_a = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 0])[:, tf.newaxis,:], [1, n_points, 1])
    p_b = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 1])[:, tf.newaxis,:], [1, n_points, 1])
    p_c = tf.tile(tf.gather(fixed_mesh_vertices, fixed_mesh_faces[:, 2])[:, tf.newaxis,:], [1, n_points, 1])
    points = tf.tile(points[tf.newaxis,: ,:], [n_faces, 1, 1])

    area_ABC = tf.abs(tf.linalg.cross(p_b -p_a, p_c - p_a)[:, :, 2])
    area_PBC = tf.abs(tf.linalg.cross(p_b -points, p_c - points)[:, :, 2])
    area_PAB = tf.abs(tf.linalg.cross(p_b -points, p_a - points)[:, :, 2])
    area_PCA = tf.abs(tf.linalg.cross(p_c -points, p_a - points)[:, :, 2])

    u= tf.divide(area_PBC, area_ABC)
    v= tf.divide(area_PCA, area_ABC)
    w= tf.divide(area_PAB, area_ABC)
    mask = u + w + v>1+ 1e-6
    u = tf.where(mask, tf.zeros_like(u), u)
    v = tf.where(mask, tf.zeros_like(v), v)
    w = tf.where(mask, tf.zeros_like(w), w)

    return v, w, u
//  [python]
def balance_classes(y_target,num_classes):
    """ Calculates the class weights needed to balance the classes, based on the number of samples of each class in the \
        batch of data.

    Args:
        y_target (tensor): One-hot classification labels (1D vector). Shape [numSamples x numClasses]
        num_classes (int):

    Returns:
        (tensor): A weighting for each class that balances their contribution to the loss. Shape [numClasses].
    """
    y_target = tf.reshape( y_target, [-1, num_classes] )
    class_count = tf.add( tf.reduce_sum( y_target, axis=0 ), tf.constant( [1]*num_classes, dtype=tf.float32 ) )
    class_weights = tf.multiply( tf.truediv( tf.ones( ( 1, num_classes) ), class_count ), tf.reduce_max( class_count ) )

    return class_weights
//  [python]
def balance_classes(y_target,num_classes):
    """ Calculates the class weights needed to balance the classes, based on the number of samples of each class in the \
        batch of data.

    Args:
        y_target (tensor): One-hot classification labels (1D vector). Shape [numSamples x numClasses]
        num_classes (int):

    Returns:
        (tensor): A weighting for each class that balances their contribution to the loss. Shape [numClasses].
    """
    y_target = tf.reshape( y_target, [-1, num_classes] )
    class_count = tf.add( tf.reduce_sum( y_target, axis=0 ), tf.constant( [1]*num_classes, dtype=tf.float32 ) )
    class_weights = tf.multiply( tf.divide( tf.ones( ( 1, num_classes) ), class_count ), tf.reduce_max( class_count ) )

    return class_weights
//  [python]
def get_losses( self, output_vectors, idx_segments, masks ):
        '''Returns the metric loss for 'num_pixels' embedding vectors. 
        
        Args:
            output_imgs: Tensor of images output by the decoder.
            desired_imgs: Tensor of target images to be output by the decoder.
            masks: Tensor of masks to be applied when computing sum of squares
                    loss.
            
        Returns:
            losses: list of tensors representing each loss component
        '''
        print('setting up losses...')
        self.output_images = output_vectors
        self.target_images = idx_segments
        self.masks = masks

        with tf.variable_scope('losses'):
            last_axis = 2
            fir, sec, seg_id = tf.unstack(idx_segments, axis=last_axis)

            idxes = tf.stack([self.batch_index_slice, fir, sec], axis=last_axis)
            self.embed = tf.gather_nd( output_vectors, idxes )
            embed = self.embed
            square = tf.reduce_sum( embed*embed, axis=-1 )
            square_t = tf.expand_dims(square, axis=-1)
            square = tf.expand_dims(square, axis=1)

            pairwise_dist = square - 2 * tf.matmul(embed, tf.transpose(embed, perm=[0,2,1])) + square_t 
            pairwise_dist = tf.clip_by_value( pairwise_dist, 0, 80)
            #pairwise_dist = 0 - pairwise_dist
            self.pairwise_dist = pairwise_dist
            pairwise_exp = tf.exp(pairwise_dist) + 1
            sigma = tf.truediv(2 ,  pairwise_exp)
            sigma = tf.clip_by_value(sigma,1e-7,1.0 - 1e-7)
            self.sigma = sigma
            same = tf.log(sigma)
            diff = tf.log(1 - sigma)

            self.same = same
            self.diff = diff
            
            seg_id_i = tf.tile(tf.expand_dims(seg_id, -1), [1, 1, self.num_pixels])
            seg_id_j = tf.transpose(seg_id_i, perm=[0,2,1])

            seg_comp = tf.equal(seg_id_i, seg_id_j)
            seg_same = tf.cast(seg_comp, self.input_type) 
            seg_diff = 1 - seg_same

            loss_matrix = seg_same * same + seg_diff * diff
            reduced_loss = 0 - tf.reduce_mean(loss_matrix) # / self.num_pixels
         
        tf.add_to_collection(tf.GraphKeys.LOSSES, reduced_loss)
        self.metric_loss = reduced_loss
        losses = [reduced_loss]
        return losses
//  [python]
def get_losses( self, output_vectors, idx_segments, masks ):
        '''Returns the metric loss for 'num_pixels' embedding vectors. 
        
        Args:
            output_imgs: Tensor of images output by the decoder.
            desired_imgs: Tensor of target images to be output by the decoder.
            masks: Tensor of masks to be applied when computing sum of squares
                    loss.
            
        Returns:
            losses: list of tensors representing each loss component
        '''
        print('setting up losses...')
        self.output_images = output_vectors
        self.target_images = idx_segments
        self.masks = masks

        with tf.variable_scope('losses'):
            last_axis = 2
            fir, sec, seg_id = tf.unstack(idx_segments, axis=last_axis)

            idxes = tf.stack([self.batch_index_slice, fir, sec], axis=last_axis)
            self.embed = tf.gather_nd( output_vectors, idxes )
            embed = self.embed
            square = tf.reduce_sum( embed*embed, axis=-1 )
            square_t = tf.expand_dims(square, axis=-1)
            square = tf.expand_dims(square, axis=1)

            pairwise_dist = square - 2 * tf.matmul(embed, tf.transpose(embed, perm=[0,2,1])) + square_t 
            pairwise_dist = tf.clip_by_value( pairwise_dist, 0, 80)
            #pairwise_dist = 0 - pairwise_dist
            self.pairwise_dist = pairwise_dist
            pairwise_exp = tf.exp(pairwise_dist) + 1
            sigma = tf.divide(2 ,  pairwise_exp)
            sigma = tf.clip_by_value(sigma,1e-7,1.0 - 1e-7)
            self.sigma = sigma
            same = tf.log(sigma)
            diff = tf.log(1 - sigma)

            self.same = same
            self.diff = diff
            
            seg_id_i = tf.tile(tf.expand_dims(seg_id, -1), [1, 1, self.num_pixels])
            seg_id_j = tf.transpose(seg_id_i, perm=[0,2,1])

            seg_comp = tf.equal(seg_id_i, seg_id_j)
            seg_same = tf.cast(seg_comp, self.input_type) 
            seg_diff = 1 - seg_same

            loss_matrix = seg_same * same + seg_diff * diff
            reduced_loss = 0 - tf.reduce_mean(loss_matrix) # / self.num_pixels
         
        tf.add_to_collection(tf.GraphKeys.LOSSES, reduced_loss)
        self.metric_loss = reduced_loss
        losses = [reduced_loss]
        return losses
//  [python]
def compute_depthmap_l1_loss(pred_batch, pred_mask_batch, pred_gpmmmask_batch, gt_batch):
    batch_size = pred_batch.shape[0]
    # l1
    curr_render_error = pred_batch - gt_batch

    curr_render_src_error = tf.abs(curr_render_error)

    curr_render_src_error = curr_render_src_error * pred_mask_batch
    curr_render_src_error = curr_render_src_error * pred_gpmmmask_batch


    error_max = tf.reduce_max(tf.reshape(curr_render_src_error, [batch_size, -1]), axis=1)
    curr_render_src_error_norm = tf.div(curr_render_src_error, tf.reshape(error_max, [batch_size, 1, 1, 1]) + 1e-6)

    curr_render_src_error_visual = pixel_error_heatmap(curr_render_src_error_norm)

    # curr_render_option_sum = pred_mask_batch * pred_gpmmmask_batch
    # curr_render_option_sum = tf.reduce_sum(curr_render_option_sum, axis=[1, 2, 3])
    #
    # curr_render_src_error = tf.reduce_sum(curr_render_src_error, axis=[1, 2, 3])
    # curr_render_src_error = curr_render_src_error / (curr_render_option_sum + 1e-6)

    curr_render_src_error = tf.reduce_mean(curr_render_src_error, axis=[1, 2, 3])

    return curr_render_src_error, curr_render_src_error_visual
//  [python]
def compute_depthmap_l1_loss(pred_batch, pred_mask_batch, pred_gpmmmask_batch, gt_batch):
    batch_size = pred_batch.shape[0]
    # l1
    curr_render_error = pred_batch - gt_batch

    curr_render_src_error = tf.abs(curr_render_error)

    curr_render_src_error = curr_render_src_error * pred_mask_batch
    curr_render_src_error = curr_render_src_error * pred_gpmmmask_batch


    error_max = tf.reduce_max(tf.reshape(curr_render_src_error, [batch_size, -1]), axis=1)
    curr_render_src_error_norm = tf.divide(curr_render_src_error, tf.reshape(error_max, [batch_size, 1, 1, 1]) + 1e-6)

    curr_render_src_error_visual = pixel_error_heatmap(curr_render_src_error_norm)

    # curr_render_option_sum = pred_mask_batch * pred_gpmmmask_batch
    # curr_render_option_sum = tf.reduce_sum(curr_render_option_sum, axis=[1, 2, 3])
    #
    # curr_render_src_error = tf.reduce_sum(curr_render_src_error, axis=[1, 2, 3])
    # curr_render_src_error = curr_render_src_error / (curr_render_option_sum + 1e-6)

    curr_render_src_error = tf.reduce_mean(curr_render_src_error, axis=[1, 2, 3])

    return curr_render_src_error, curr_render_src_error_visual
//  [python]
def norm_boxes_graph(boxes, #像素坐标(y1, x1, y2, x2)
                     shape):#像素边长(height, width)
    h, w = tf.split(tf.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.truediv(boxes - shift, scale)
//  [python]
def norm_boxes_graph(boxes, #像素坐标(y1, x1, y2, x2)
                     shape):#像素边长(height, width)
    h, w = tf.split(tf.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.divide(boxes - shift, scale)
//  [python]
def _make_placeholders(self):
        super(GaussianRBM, self)._make_placeholders()
        with tf.name_scope('input_data'):
            # divide by resp. sigmas before any operation
            self._sigma = tf.Variable(self._sigma_tmp, dtype=self._tf_dtype, name='sigma')
            self._sigma = tf.reshape(self._sigma, [1, self.n_visible])
            self._X_batch = tf.truediv(self._X_batch, self._sigma)
//  [python]
def _make_placeholders(self):
        super(GaussianRBM, self)._make_placeholders()
        with tf.name_scope('input_data'):
            # divide by resp. sigmas before any operation
            self._sigma = tf.Variable(self._sigma_tmp, dtype=self._tf_dtype, name='sigma')
            self._sigma = tf.reshape(self._sigma, [1, self.n_visible])
            self._X_batch = tf.divide(self._X_batch, self._sigma)
//  [python]
def compute_ncc_loss(a, b):
    vector_a = slim.flatten(a)[0]
    vector_b = slim.flatten(b)[0]
    mean_a, var_a = tf.nn.moments(vector_a,axes=0)
    mean_b, var_b = tf.nn.moments(vector_b,axes=0)
    new_a = tf.truediv((vector_a-mean_a),tf.sqrt(var_a)+1e-7)
    new_b = tf.truediv((vector_b-mean_b),tf.sqrt(var_b)+1e-7)
    return tf.abs(tf.reduce_mean(new_a*new_b))
//  [python]
def compute_ncc_loss(a, b):
    vector_a = slim.flatten(a)[0]
    vector_b = slim.flatten(b)[0]
    mean_a, var_a = tf.nn.moments(vector_a,axes=0)
    mean_b, var_b = tf.nn.moments(vector_b,axes=0)
    new_a = tf.divide((vector_a-mean_a),tf.sqrt(var_a)+1e-7)
    new_b = tf.divide((vector_b-mean_b),tf.sqrt(var_b)+1e-7)
    return tf.abs(tf.reduce_mean(new_a*new_b))
//  [python]
def tf_box_8c_to_offsets(boxes_8c,
                         boxes_8c_gt):
    """Converts corner boxes to corner offsets.
    It subtracts the ground-truth box corners from the predicted corners
    and normalizes the offsets by the diagonal of the proposed boxes.

    Args:
        boxes_8c: A tensor of shape (N x 3 x 8) representing the box corners.
        boxes_8c_gt: A tensor of shape (N x 3 x 8) representing the box
            corners ground-truth.
    Returns:
        A tensor of dim (N x 3 x 8) representing the offsets.
    """

    # Get the diagonal of the boxes
    diagonals = tf_box_8c_diagonal_length(boxes_8c)
    offsets = tf.subtract(boxes_8c_gt, boxes_8c)

    # Reshape the offsets to a (24 x N) vector
    reshaped_offsets = tf.reshape(offsets, (24, -1))
    ones = tf.ones_like(reshaped_offsets)
    # This gives diagonals of shape (24 x N)
    # This now enables us to divide acorss N batches
    diagonals_mult = tf.multiply(ones, diagonals)
    # Normalize the offsets by the box_8c diagonal
    offsets_norm = tf.truediv(reshaped_offsets, diagonals_mult)

    reshaped_offsets_norm = tf.reshape(offsets_norm,
                                       [-1, 3, 8])

    return reshaped_offsets_norm
//  [python]
def tf_box_8c_to_offsets(boxes_8c,
                         boxes_8c_gt):
    """Converts corner boxes to corner offsets.
    It subtracts the ground-truth box corners from the predicted corners
    and normalizes the offsets by the diagonal of the proposed boxes.

    Args:
        boxes_8c: A tensor of shape (N x 3 x 8) representing the box corners.
        boxes_8c_gt: A tensor of shape (N x 3 x 8) representing the box
            corners ground-truth.
    Returns:
        A tensor of dim (N x 3 x 8) representing the offsets.
    """

    # Get the diagonal of the boxes
    diagonals = tf_box_8c_diagonal_length(boxes_8c)
    offsets = tf.subtract(boxes_8c_gt, boxes_8c)

    # Reshape the offsets to a (24 x N) vector
    reshaped_offsets = tf.reshape(offsets, (24, -1))
    ones = tf.ones_like(reshaped_offsets)
    # This gives diagonals of shape (24 x N)
    # This now enables us to divide acorss N batches
    diagonals_mult = tf.multiply(ones, diagonals)
    # Normalize the offsets by the box_8c diagonal
    offsets_norm = tf.divide(reshaped_offsets, diagonals_mult)

    reshaped_offsets_norm = tf.reshape(offsets_norm,
                                       [-1, 3, 8])

    return reshaped_offsets_norm
//  [python]
def loss_function_reconstruction_1D(y_reconstructed,y_target,func='SSE'):
    """ Reconstruction loss function op, comparing 1D tensors for network reconstruction and target.

    Args:
        y_reconstructed (tensor): Output of network (reconstructed 1D vector). Shape [numSamples x inputSize].
        y_target (tensor): What the network is trying to reconstruct (1D vector). Shape [numSamples x inputSize].
        func (string): The name of the loss function to be used. 'SSE'-sum of square errors,'CSA'-cosine spectral angle, \
            'SA'-spectral angle, 'SID'-spectral information divergence.

    Returns:
        (tensor): Reconstruction loss.
    """
    if func == 'SSE':
        # sum of squared errors loss
        loss = tf.reduce_sum( tf.square(y_target - y_reconstructed) )

    elif func == 'CSA':
        # cosine of spectral angle loss
        normalize_r = tf.math.l2_normalize(tf.transpose(y_reconstructed),axis=0)
        normalize_t = tf.math.l2_normalize(tf.transpose(y_target),axis=0)
        loss = tf.reduce_sum( 1 - tf.reduce_sum(tf.multiply(normalize_r, normalize_t),axis=0 ) )

    elif func == 'SA':
        # spectral angle loss
        normalize_r = tf.math.l2_normalize(tf.transpose(y_reconstructed),axis=0)
        normalize_t = tf.math.l2_normalize(tf.transpose(y_target),axis=0)
        loss = tf.reduce_sum( tf.math.acos(tf.reduce_sum(tf.multiply(normalize_r, normalize_t),axis=0 ) ) )

    elif func == 'SID':
        # spectral information divergence loss
        t = tf.transpose(y_target) / tf.reduce_sum(tf.transpose(y_target),axis=0)
        r = tf.transpose(y_reconstructed) / tf.reduce_sum(tf.transpose(y_reconstructed),axis=0)
        loss = tf.reduce_sum( tf.reduce_sum( tf.multiply(t,tf.math.log(t/r))) , axis=0) + tf.reduce_sum( tf.reduce_sum( tf.multiply(r,tf.math.log(r/t))) , axis=0)
    else:
        raise ValueError('unknown loss function: %s. Use SSE, CSA, SA or SID.' % func)

    return loss
//  [python]
def loss_function_reconstruction_1D(y_reconstructed,y_target,func='SSE'):
    """ Reconstruction loss function op, comparing 1D tensors for network reconstruction and target.

    Args:
        y_reconstructed (tensor): Output of network (reconstructed 1D vector). Shape [numSamples x inputSize].
        y_target (tensor): What the network is trying to reconstruct (1D vector). Shape [numSamples x inputSize].
        func (string): The name of the loss function to be used. 'SSE'-sum of square errors,'CSA'-cosine spectral angle, \
            'SA'-spectral angle, 'SID'-spectral information divergence.

    Returns:
        (tensor): Reconstruction loss.
    """
    if func == 'SSE':
        # sum of squared errors loss
        loss = tf.reduce_sum( tf.square(y_target - y_reconstructed) )

    elif func == 'CSA':
        # cosine of spectral angle loss
        normalize_r = tf.math.l2_normalize(tf.transpose(y_reconstructed),axis=0)
        normalize_t = tf.math.l2_normalize(tf.transpose(y_target),axis=0)
        loss = tf.reduce_sum( 1 - tf.reduce_sum(tf.multiply(normalize_r, normalize_t),axis=0 ) )

    elif func == 'SA':
        # spectral angle loss
        normalize_r = tf.math.l2_normalize(tf.transpose(y_reconstructed),axis=0)
        normalize_t = tf.math.l2_normalize(tf.transpose(y_target),axis=0)
        loss = tf.reduce_sum( tf.math.acos(tf.reduce_sum(tf.multiply(normalize_r, normalize_t),axis=0 ) ) )

    elif func == 'SID':
        # spectral information divergence loss
        t = tf.divide( tf.transpose(y_target) , tf.reduce_sum(tf.transpose(y_target),axis=0) )
        r = tf.divide( tf.transpose(y_reconstructed) , tf.reduce_sum(tf.transpose(y_reconstructed),axis=0) )
        loss = tf.reduce_sum( tf.reduce_sum( tf.multiply(t,tf.log(tf.divide(t,r))) , axis=0)
                              + tf.reduce_sum( tf.multiply(r,tf.log(tf.divide(r,t))) , axis=0) )
    else:
        raise ValueError('unknown loss function: %s. Use SSE, CSA, SA or SID.' % func)

    return loss
//  [python]
def _free_energy(self, v):
        with tf.name_scope('free_energy'):
            T1 = tf.div(tf.reshape(self._vb, [1, self.n_visible]), self._sigma)
            T2 = tf.square(tf.subtract(v, T1))
            T3 = 0.5 * tf.reduce_sum(T2, axis=1)
            T4 = -tf.reduce_sum(tf.nn.softplus(self._propup(v) + self._hb), axis=1)
            fe = tf.reduce_mean(T3 + T4, axis=0)
        return fe
//  [python]
def _free_energy(self, v):
        with tf.name_scope('free_energy'):
            T1 = tf.divide(tf.reshape(self._vb, [1, self.n_visible]), self._sigma)
            T2 = tf.square(tf.subtract(v, T1))
            T3 = 0.5 * tf.reduce_sum(T2, axis=1)
            T4 = -tf.reduce_sum(tf.nn.softplus(self._propup(v) + self._hb), axis=1)
            fe = tf.reduce_mean(T3 + T4, axis=0)
        return fe
//  [python]
def _hard_sigmoid(x, name=None):
    return tf.div(tf.nn.relu6(x + 3.), 6., name=name)
//  [python]
def _hard_sigmoid(x, name=None):
    return tf.divide(tf.nn.relu6(x + 3.), 6., name=name)
//  [python]
def l2_normalize_channel_wise(features):
  norms = tf.norm(features, ord='euclidean', axis=3)
  # expanding the norms tensor to support broadcast division
  norms_expanded = tf.expand_dims(norms, axis=3)
  features_norm = tf.truediv(features, norms_expanded)
  return features_norm
//  [python]
def l2_normalize_channel_wise(features):
  norms = tf.norm(features, ord='euclidean', axis=3)
  # expanding the norms tensor to support broadcast division
  norms_expanded = tf.expand_dims(norms, axis=3)
  features_norm = tf.divide(features, norms_expanded)
  return features_norm
//  [python]
def srn(x, depth_radius, alpha=1.0, beta=0.5, name=None):
    squared_sum = tf.nn.depthwise_conv2d(
        tf.square(x),
        tf.ones([depth_radius] * 2 + [tf.shape(x)[3], 1], dtype=tf.float32),
        [1, 1, 1, 1],
        'SAME')
    alpha = tf.constant(alpha / (depth_radius ** 2), dtype=tf.float32)
    beta = tf.constant(beta, dtype=tf.float32)
    return tf.truediv(x, (1.0 + alpha * squared_sum) ** beta, name=name)
//  [python]
def srn(x, depth_radius, alpha=1.0, beta=0.5, name=None):
    # Refer to the following code snippet
    # https://github.com/tensorflow/tensorflow/issues/1246#issuecomment-188588051
    squared_sum = tf.nn.depthwise_conv2d(
        tf.square(x),
        tf.ones([depth_radius] * 2 + [tf.shape(x)[3], 1], dtype=tf.float32),
        [1, 1, 1, 1],
        'SAME')
    alpha = tf.constant(alpha / (depth_radius ** 2), dtype=tf.float32)
    beta = tf.constant(beta, dtype=tf.float32)
    return tf.divide(x, (1.0 + alpha * squared_sum) ** beta, name=name)
//  [python]
def _evolve_trotter_gates(psi,
                          layers,
                          step_size,
                          num_steps,
                          euclidean=False,
                          callback=None):
  """Evolve an initial wavefunction psi via gates specified in `layers`.

  If the evolution is euclidean, the wavefunction will be normalized
  after each step.
  """
  t = 0.0
  for i in range(num_steps):
    psi = apply_circuit(psi, layers)
    if euclidean:
      psi = psi / tf.norm(psi)
    t += step_size
    if callback is not None:
      callback(psi, t, i)

  return psi, t
//  [python]
def _evolve_trotter_gates(psi,
                          layers,
                          step_size,
                          num_steps,
                          euclidean=False,
                          callback=None):
  """Evolve an initial wavefunction psi via gates specified in `layers`.

  If the evolution is euclidean, the wavefunction will be normalized
  after each step.
  """
  t = 0.0
  for i in range(num_steps):
    psi = apply_circuit(psi, layers)
    if euclidean:
      psi = tf.divide(psi, tf.norm(psi))
    t += step_size
    if callback is not None:
      callback(psi, t, i)

  return psi, t
//  [python]
def sum_normalize(cs, axis=TensorAxis.C):
        reduce_sum = tf.reduce_sum(cs, axis, keep_dims=True, name='sum')
        return tf.truediv(cs, reduce_sum, name='sumNormalized')
//  [python]
def sum_normalize(cs, axis=TensorAxis.C):
        reduce_sum = tf.reduce_sum(cs, axis, keep_dims=True, name='sum')
        return tf.divide(cs, reduce_sum, name='sumNormalized')
//  [python]
def create_discrete_action_masking_layer(all_logits, action_masks, action_size):
        """
        Creates a masking layer for the discrete actions
        :param all_logits: The concatenated unnormalized action probabilities for all branches
        :param action_masks: The mask for the logits. Must be of dimension [None x total_number_of_action]
        :param action_size: A list containing the number of possible actions for each branch
        :return: The action output dimension [batch_size, num_branches] and the concatenated normalized logits
        """
        action_idx = [0] + list(np.cumsum(action_size))
        branches_logits = [all_logits[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
        branch_masks = [action_masks[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
        raw_probs = [tf.multiply(tf.nn.softmax(branches_logits[k]) + 1.0e-10, branch_masks[k])
                     for k in range(len(action_size))]
        normalized_probs = [
            tf.truediv(raw_probs[k], tf.reduce_sum(raw_probs[k], axis=1, keepdims=True))
            for k in range(len(action_size))]
        output = tf.concat([tf.multinomial(tf.log(normalized_probs[k]), 1) for k in range(len(action_size))], axis=1)
        return output, tf.concat([tf.log(normalized_probs[k] + 1.0e-10) for k in range(len(action_size))], axis=1)
//  [python]
def create_discrete_action_masking_layer(all_logits, action_masks, action_size):
        """
        Creates a masking layer for the discrete actions
        :param all_logits: The concatenated unnormalized action probabilities for all branches
        :param action_masks: The mask for the logits. Must be of dimension [None x total_number_of_action]
        :param action_size: A list containing the number of possible actions for each branch
        :return: The action output dimension [batch_size, num_branches] and the concatenated normalized logits
        """
        action_idx = [0] + list(np.cumsum(action_size))
        branches_logits = [all_logits[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
        branch_masks = [action_masks[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
        raw_probs = [tf.multiply(tf.nn.softmax(branches_logits[k]) + 1.0e-10, branch_masks[k])
                     for k in range(len(action_size))]
        normalized_probs = [
            tf.divide(raw_probs[k], tf.reduce_sum(raw_probs[k], axis=1, keepdims=True))
            for k in range(len(action_size))]
        output = tf.concat([tf.multinomial(tf.log(normalized_probs[k]), 1) for k in range(len(action_size))], axis=1)
        return output, tf.concat([tf.log(normalized_probs[k] + 1.0e-10) for k in range(len(action_size))], axis=1)
//  [python]
def l2_normalize_channelwise(features):
        norms = tf.norm(features, ord='euclidean', axis=TensorAxis.C, name='norm')
        # expanding the norms tensor to support broadcast division
        norms_expanded = tf.expand_dims(norms, TensorAxis.C)
        features = tf.truediv(features, norms_expanded, name='normalized')
        return features
//  [python]
def l2_normalize_channelwise(features):
        norms = tf.norm(features, ord='euclidean', axis=TensorAxis.C, name='norm')
        # expanding the norms tensor to support broadcast division
        norms_expanded = tf.expand_dims(norms, TensorAxis.C)
        features = tf.divide(features, norms_expanded, name='normalized')
        return features
//  [python]
def scale_image_values(tf_tensor, name=None):
    return tf.truediv(tf_tensor, 255.0, name=name)
//  [python]
def scale_image_values(tf_tensor, name=None):
    return tf.divide(tf_tensor, 255.0, name=name)
//  [python]
def vgg_loss(feat_net, feat_weights, n_layers, reg=0.1):
    def loss_fcn(y_true, y_pred):
        y_true_feat = feat_net(Lambda(vgg_preprocess)(y_true))
        y_pred_feat = feat_net(Lambda(vgg_preprocess)(y_pred))

        loss = []
        for j in range(n_layers):

            std = feat_weights[str(j)][1] + reg
            std = tf.expand_dims(tf.expand_dims(tf.expand_dims(std, 0), 0), 0)
            d = tf.subtract(y_true_feat[j], y_pred_feat[j])
            loss_j = tf.reduce_mean(tf.abs(tf.truediv(d, std)))

            if j == 0:
                loss = loss_j
            else:
                loss = tf.add(loss, loss_j)
        return loss / (n_layers * 1.0)

    return loss_fcn
//  [python]
def vgg_loss(feat_net, feat_weights, n_layers, reg=0.1):
    def loss_fcn(y_true, y_pred):
        y_true_feat = feat_net(Lambda(vgg_preprocess)(y_true))
        y_pred_feat = feat_net(Lambda(vgg_preprocess)(y_pred))

        loss = []
        for j in range(n_layers):

            std = feat_weights[str(j)][1] + reg
            std = tf.expand_dims(tf.expand_dims(tf.expand_dims(std, 0), 0), 0)
            d = tf.subtract(y_true_feat[j], y_pred_feat[j])
            loss_j = tf.reduce_mean(tf.abs(tf.divide(d, std)))

            if j == 0:
                loss = loss_j
            else:
                loss = tf.add(loss, loss_j)
        return loss / (n_layers * 1.0)

    return loss_fcn
//  [python]
def normalization(tensor):
    #normalized_ds = data.map(lambda x: normalization_layer(x))
    tensor = tf.image.resize(
    tensor, (128,128))
    tensor = tf.subtract(tf.truediv(tensor, 127.5), 1)
    return tensor
//  [python]
def normalization(tensor):
    #normalized_ds = data.map(lambda x: normalization_layer(x))
    tensor = tf.image.resize(
    tensor, (128,128))
    tensor = tf.subtract(tf.divide(tensor, 127.5), 1)
    return tensor
//  [python]
def _create_variable_(name, initializer, fan_size, trainable=True, weight_scale=1.0):
    """ This function pins variables to cpu

    tf.get_variable is used instead of tf.Variable, so that variable with the same name will not be re-initialized
    """
    # define initialization method
    if initializer == 'zeros':
        initializer_fun = tf.zeros(fan_size)
    elif initializer == 'ones':
        initializer_fun = tf.multiply(tf.ones(fan_size), weight_scale)
    elif initializer == 'xavier':
        initializer_fun = xavier_init(fan_size[0], fan_size[1], weight_scale=weight_scale)
    elif initializer == 'normal_in':
        initializer_fun = tf.random_normal(fan_size, mean=0.0, stddev=tf.div(1.0, tf.sqrt(fan_size[0])),
                                           dtype=tf.float32)
    else:
        raise AttributeError('Initialization method not supported.')

    return tf.get_variable(name=name, initializer=initializer_fun, trainable=trainable)
//  [python]
def _create_variable_(name, initializer, fan_size, trainable=True, weight_scale=1.0):
    """ This function pins variables to cpu

    tf.get_variable is used instead of tf.Variable, so that variable with the same name will not be re-initialized
    """
    # define initialization method
    if initializer == 'zeros':
        initializer_fun = tf.zeros(fan_size)
    elif initializer == 'ones':
        initializer_fun = tf.multiply(tf.ones(fan_size), weight_scale)
    elif initializer == 'xavier':
        initializer_fun = xavier_init(fan_size[0], fan_size[1], weight_scale=weight_scale)
    elif initializer == 'normal_in':
        initializer_fun = tf.random_normal(fan_size, mean=0.0, stddev=tf.divide(1.0, tf.sqrt(fan_size[0])),
                                           dtype=tf.float32)
    else:
        raise AttributeError('Initialization method not supported.')

    return tf.get_variable(name=name, initializer=initializer_fun, trainable=trainable)
//  [python]
def norm_boxes_graph(boxes, shape):
    """Converts boxes from pixel coordinates to normalized coordinates.
    boxes: [..., (y1, x1, y2, x2)] in pixel coordinates
    shape: [..., (height, width)] in pixels
    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.
    Returns:
        [..., (y1, x1, y2, x2)] in normalized coordinates
    """
    h, w = tf.split(tf.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.truediv(boxes - shift, scale)
//  [python]
def norm_boxes_graph(boxes, shape):
    """Converts boxes from pixel coordinates to normalized coordinates.
    boxes: [..., (y1, x1, y2, x2)] in pixel coordinates
    shape: [..., (height, width)] in pixels
    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.
    Returns:
        [..., (y1, x1, y2, x2)] in normalized coordinates
    """
    h, w = tf.split(tf.cast(shape, tf.float32), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.divide(boxes - shift, scale)
//  [python]
def loss_l2(self, var_list):
    for c in self.weight_decay_constraint:
      var_list = [v for v in var_list if c not in v.name]
    loss_l2 = tf.add_n([tf.nn.l2_loss(v) for v in var_list])
    return tf.truediv(loss_l2, self.strategy.num_replicas_in_sync)
//  [python]
def loss_l2(self, var_list):
    for c in self.weight_decay_constraint:
      var_list = [v for v in var_list if c not in v.name]
    loss_l2 = tf.add_n([tf.nn.l2_loss(v) for v in var_list])
    return tf.divide(loss_l2, self.strategy.num_replicas_in_sync)
//  [python]
def replace_weights(self, new_weights):
        for e, weights in enumerate(self.weights):
            if isinstance(weights, str):
                weights = tf.get_default_graph().get_tensor_by_name(weights)
            if weights.name.endswith('/batch_normalization/beta:0'):
                # Batch norm beta is written as bias before the batch
                # normalization in the weight file for backwards
                # compatibility reasons.
                bias = tf.constant(new_weights[e], shape=weights.shape)
                # Weight file order: bias, means, variances
                var = tf.constant(new_weights[e + 2], shape=weights.shape)
                new_beta = tf.div(bias, tf.sqrt(var + tf.constant(1e-5)))
                self.assign(weights, new_beta)
            elif weights.shape.ndims == 4:
                # Convolution weights need a transpose
                #
                # TF (kYXInputOutput)
                # [filter_height, filter_width, in_channels, out_channels]
                #
                # Leela/cuDNN/Caffe (kOutputInputYX)
                # [output, input, filter_size, filter_size]
                s = weights.shape.as_list()
                shape = [s[i] for i in [3, 2, 0, 1]]
                new_weight = tf.constant(new_weights[e], shape=shape)
                self.assign(weights, tf.transpose(new_weight, [2, 3, 1, 0]))
            elif weights.shape.ndims == 2:
                # Fully connected layers are [in, out] in TF
                #
                # [out, in] in Leela
                #
                s = weights.shape.as_list()
                shape = [s[i] for i in [1, 0]]
                new_weight = tf.constant(new_weights[e], shape=shape)
                self.assign(weights, tf.transpose(new_weight, [1, 0]))
            else:
                # Biases, batchnorm etc
                new_weight = tf.constant(new_weights[e], shape=weights.shape)
                self.assign(weights, new_weight)
//  [python]
def replace_weights(self, new_weights):
        for e, weights in enumerate(self.weights):
            if isinstance(weights, str):
                weights = tf.get_default_graph().get_tensor_by_name(weights)
            if weights.name.endswith('/batch_normalization/beta:0'):
                # Batch norm beta is written as bias before the batch
                # normalization in the weight file for backwards
                # compatibility reasons.
                bias = tf.constant(new_weights[e], shape=weights.shape)
                # Weight file order: bias, means, variances
                var = tf.constant(new_weights[e + 2], shape=weights.shape)
                new_beta = tf.divide(bias, tf.sqrt(var + tf.constant(1e-5)))
                self.assign(weights, new_beta)
            elif weights.shape.ndims == 4:
                # Convolution weights need a transpose
                #
                # TF (kYXInputOutput)
                # [filter_height, filter_width, in_channels, out_channels]
                #
                # Leela/cuDNN/Caffe (kOutputInputYX)
                # [output, input, filter_size, filter_size]
                s = weights.shape.as_list()
                shape = [s[i] for i in [3, 2, 0, 1]]
                new_weight = tf.constant(new_weights[e], shape=shape)
                self.assign(weights, tf.transpose(new_weight, [2, 3, 1, 0]))
            elif weights.shape.ndims == 2:
                # Fully connected layers are [in, out] in TF
                #
                # [out, in] in Leela
                #
                s = weights.shape.as_list()
                shape = [s[i] for i in [1, 0]]
                new_weight = tf.constant(new_weights[e], shape=shape)
                self.assign(weights, tf.transpose(new_weight, [1, 0]))
            else:
                # Biases, batchnorm etc
                new_weight = tf.constant(new_weights[e], shape=weights.shape)
                self.assign(weights, new_weight)
//  [python]
def train_step(self, iterator):
    """Executes each train step."""

    def step_fn(data):
      replica_context = tf.distribute.get_replica_context()
      xo, xc = data[0], data[1]
      x = tf.concat((xo, xc), axis=0)
      y = tf.concat((tf.zeros(
          xo.shape[0], dtype=tf.int32), tf.ones(xc.shape[0], dtype=tf.int32)),
                    axis=0)
      with tf.GradientTape() as tape:
        logits = self.model(x, training=True)['logits']
        loss_xe = tf.keras.losses.sparse_categorical_crossentropy(
            y, logits, from_logits=True)
        loss_xe = tf.reduce_sum(loss_xe) / self.cross_replica_concat(loss_xe,
                                      replica_context=replica_context).shape[0]
        loss_l2 = self.loss_l2(self.model.trainable_weights)
        loss = loss_xe + self.weight_decay * loss_l2
      grad = tape.gradient(loss, self.model.trainable_weights)
      self.optimizer.apply_gradients(zip(grad, self.model.trainable_weights))
      # monitor
      self.metrics['loss.train'].update_state(loss)
      self.metrics['loss.xe'].update_state(loss_xe)
      self.metrics['loss.L2'].update_state(loss_l2)
      self.metrics['acc.train'].update_state(y, tf.argmax(logits, axis=1))

    # Call one step
    self.strategy.run(step_fn, args=(next(iterator),))
//  [python]
def train_step(self, iterator):
    """Executes each train step."""

    def step_fn(data):
      replica_context = tf.distribute.get_replica_context()
      xo, xc = data[0], data[1]
      x = tf.concat((xo, xc), axis=0)
      y = tf.concat((tf.zeros(
          xo.shape[0], dtype=tf.int32), tf.ones(xc.shape[0], dtype=tf.int32)),
                    axis=0)
      with tf.GradientTape() as tape:
        logits = self.model(x, training=True)['logits']
        loss_xe = tf.keras.losses.sparse_categorical_crossentropy(
            y, logits, from_logits=True)
        loss_xe = tf.divide(
            tf.reduce_sum(loss_xe),
            self.cross_replica_concat(loss_xe,
                                      replica_context=replica_context).shape[0])
        loss_l2 = self.loss_l2(self.model.trainable_weights)
        loss = loss_xe + self.weight_decay * loss_l2
      grad = tape.gradient(loss, self.model.trainable_weights)
      self.optimizer.apply_gradients(zip(grad, self.model.trainable_weights))
      # monitor
      self.metrics['loss.train'].update_state(loss)
      self.metrics['loss.xe'].update_state(loss_xe)
      self.metrics['loss.L2'].update_state(loss_l2)
      self.metrics['acc.train'].update_state(y, tf.argmax(logits, axis=1))

    # Call one step
    self.strategy.run(step_fn, args=(next(iterator),))
//  [python]
def local_response_normalization(x, scope_prefix='', data_format='channels_first'):
    """ This function applies a variant of local response normalization to the activations.
    It normalizes x by the sqrt of squared sum along channel dimension.

    The local response normalization is used in following paper:
    Karras T., Aila T., Laine S., Lehtinen J. (2017).
    Progressive Growing of GANs for Improved Quality, Stability, And Variation,

    :param x:
    :param scope_prefix:
    :param data_format:
    :return:
    """
    with tf.name_scope(scope_prefix + 'LRN'):
        axis = 1 if data_format in ['channels_first', 'NCHW'] else -1
        return x / tf.sqrt(tf.reduce_mean(tf.square(x), axis=axis, keepdims=True) + FLAGS.EPSI)
//  [python]
def local_response_normalization(x, scope_prefix='', data_format='channels_first'):
    """ This function applies a variant of local response normalization to the activations.
    It normalizes x by the sqrt of squared sum along channel dimension.

    The local response normalization is used in following paper:
    Karras T., Aila T., Laine S., Lehtinen J. (2017).
    Progressive Growing of GANs for Improved Quality, Stability, And Variation,

    :param x:
    :param scope_prefix:
    :param data_format:
    :return:
    """
    with tf.name_scope(scope_prefix + 'LRN'):
        axis = 1 if data_format in ['channels_first', 'NCHW'] else -1
        return tf.divide(x, tf.sqrt(tf.reduce_mean(tf.square(x), axis=axis, keepdims=True) + FLAGS.EPSI))
//  [python]
def call(self, inputs, mask=None):
        # Import graph tensors
        # scalar_features = (samples, max_atoms, atom_feat)
        # vector_features = (samples, max_atoms, coor_dims, atom_feat)
        scalar_features, vector_features = inputs

        # Get parameters
        coor_dims = int(vector_features.shape[2])
        atom_feat = int(vector_features.shape[-1])

        # Integrate over atom axis
        if self.pooling == "sum":
            scalar_features = tf.reduce_sum(scalar_features, axis=1)
            vector_features = tf.reduce_sum(vector_features, axis=1)

        elif self.pooling == "max":
            scalar_features = tf.reduce_max(scalar_features, axis=1)

            vector_features = tf.transpose(vector_features, perm=[0, 2, 3, 1])
            size = tf.sqrt(tf.reduce_sum(tf.square(vector_features), axis=1))
            idx = tf.reshape(tf.argmax(size, axis=-1, output_type=tf.int32), [-1, 1, atom_feat, 1])
            idx = tf.tile(idx, [1, coor_dims, 1, 1])
            vector_features = tf.reshape(tf.batch_gather(vector_features, idx), [-1, coor_dims, atom_feat])

        # Activation
        scalar_features = self.activation(scalar_features)
        vector_features = self.activation(vector_features)

        if self.system == "spherical":
            x, y, z = tf.unstack(vector_features, axis=1)
            r = tf.sqrt(tf.square(x) + tf.square(y) + tf.square(z))
            t = tf.acos(tf.truediv(z, r + tf.cast(tf.equal(r, 0), dtype=float)))
            p = tf.atan(tf.truediv(y, x + tf.cast(tf.equal(x, 0), dtype=float)))
            vector_features = tf.stack([r, t, p], axis=1)

        return [scalar_features, vector_features]
//  [python]
def call(self, inputs, mask=None):
        # Import graph tensors
        # scalar_features = (samples, max_atoms, atom_feat)
        # vector_features = (samples, max_atoms, coor_dims, atom_feat)
        scalar_features, vector_features = inputs

        # Get parameters
        coor_dims = int(vector_features.shape[2])
        atom_feat = int(vector_features.shape[-1])

        # Integrate over atom axis
        if self.pooling == "sum":
            scalar_features = tf.reduce_sum(scalar_features, axis=1)
            vector_features = tf.reduce_sum(vector_features, axis=1)

        elif self.pooling == "max":
            scalar_features = tf.reduce_max(scalar_features, axis=1)

            vector_features = tf.transpose(vector_features, perm=[0, 2, 3, 1])
            size = tf.sqrt(tf.reduce_sum(tf.square(vector_features), axis=1))
            idx = tf.reshape(tf.argmax(size, axis=-1, output_type=tf.int32), [-1, 1, atom_feat, 1])
            idx = tf.tile(idx, [1, coor_dims, 1, 1])
            vector_features = tf.reshape(tf.batch_gather(vector_features, idx), [-1, coor_dims, atom_feat])

        # Activation
        scalar_features = self.activation(scalar_features)
        vector_features = self.activation(vector_features)

        if self.system == "spherical":
            x, y, z = tf.unstack(vector_features, axis=1)
            r = tf.sqrt(tf.square(x) + tf.square(y) + tf.square(z))
            t = tf.acos(tf.divide(z, r + tf.cast(tf.equal(r, 0), dtype=float)))
            p = tf.atan(tf.divide(y, x + tf.cast(tf.equal(x, 0), dtype=float)))
            vector_features = tf.stack([r, t, p], axis=1)

        return [scalar_features, vector_features]
//  [python]
def batch_norm(
    x: Union[tf.Tensor, tf.Variable],
    mean: Union[tf.Tensor, tf.Variable],
    variance: Union[tf.Tensor, tf.Variable],
    /,
    *,
    scale: Optional[Union[tf.Tensor, tf.Variable]] = None,
    offset: Optional[Union[tf.Tensor, tf.Variable]] = None,
    training: Optional[bool] = False,
    eps: Optional[float] = 1e-5,
    momentum: Optional[float] = 1e-1,
    data_format: Optional[str] = "NSC",
    out: Optional[
        Tuple[
            Union[tf.Tensor, tf.Variable],
            Union[tf.Tensor, tf.Variable],
            Union[tf.Tensor, tf.Variable],
        ]
    ] = None,
) -> Tuple[
    Union[tf.Tensor, tf.Variable],
    Union[tf.Tensor, tf.Variable],
    Union[tf.Tensor, tf.Variable],
]:
    xdims = len(x.shape)
    if data_format == "NCS":
        x = tf.transpose(x, perm=(0, *range(2, xdims), 1))

    runningmean = mean
    runningvariance = variance
    if training:
        n = tf.size(x) if xdims == 1 else tf.div(tf.size(x), tf.shape(x)[-1])
        n = tf.cast(n, x.dtype) if n.dtype != x.dtype else n
        dims = (0, *range(1, xdims - 1))
        mean = tf.math.reduce_mean(x, axis=dims)
        variance = tf.math.reduce_variance(x, axis=dims)
        runningmean = (1 - momentum) * runningmean + momentum * mean
        runningvariance = (1 - momentum) * runningvariance + momentum * variance * n / (
            n - 1
        )

    inv = 1.0 / tf.math.sqrt(variance + eps)
    offset = 0 if offset is None else offset
    if scale is not None:
        inv = tf.math.multiply(inv, scale)
    xnormalized = tf.math.add(tf.math.multiply(x, inv), offset)
    xnormalized = tf.math.subtract(xnormalized, tf.math.multiply(mean, inv))
    # the above approach is faster than tf.nn.batch_normalization

    if data_format == "NCS":
        xnormalized = tf.transpose(
            xnormalized, perm=(0, xdims - 1, *range(1, xdims - 1))
        )

    return xnormalized, runningmean, runningvariance