jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def depth2pts_outside(ray_o, ray_d, depth):
    '''
    ray_o, ray_d: [..., 3]
    depth: [...]; inverse of distance to sphere origin
    '''
    # note: d1 becomes negative if this mid point is behind camera
    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
    p_mid = ray_o + d1.unsqueeze(-1) * ray_d
    p_mid_norm = torch.linalg.norm(p_mid, dim=-1)
    ray_d_cos = 1. / torch.linalg.norm(ray_d, dim=-1)
    d2 = torch.sqrt(1. - p_mid_norm * p_mid_norm) * ray_d_cos
    p_sphere = ray_o + (d1 + d2).unsqueeze(-1) * ray_d

    rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
    rot_axis = rot_axis / torch.linalg.norm(rot_axis, dim=-1, keepdim=True)
    phi = torch.asin(p_mid_norm)
    theta = torch.asin(p_mid_norm * depth)  # depth is inside [0, 1]
    rot_angle = (phi - theta).unsqueeze(-1)     # [..., 1]

    # now rotate p_sphere
    # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
    p_sphere_new = p_sphere * torch.cos(rot_angle) + \
                torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
                rot_axis * torch.sum(rot_axis*p_sphere, dim=-1, keepdim=True) * (1.-torch.cos(rot_angle))
    p_sphere_new = p_sphere_new / torch.linalg.norm(p_sphere_new, dim=-1, keepdim=True)
    pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)

    # now calculate conventional depth
    depth_real = 1. / (depth + TINY_NUMBER) * torch.cos(theta) * ray_d_cos + d1
    return pts, depth_real
//  [python]
def compute_weight(self, module, do_power_iteration):
        weight = getattr(module, self.name + '_orig')
        u = getattr(module, self.name + '_u')
        v = getattr(module, self.name + '_v')
        curr_shape = weight.shape
        weight_mat = self.reshape_weight_to_matrix(weight)

        if do_power_iteration:
            with torch.no_grad():
                for _ in range(self.n_power_iterations):
                    # Spectral norm of weight equals to `u^T W v`,
                    # where `u` and `v`
                    # are the first left and right singular vectors.
                    # This power iteration produces approximations
                    # of `u` and `v`.
                    v = normalize(torch.mv(weight_mat.t(), u),
                                  dim=0, eps=self.eps, out=v)
                    u = normalize(torch.mv(weight_mat, v),
                                  dim=0, eps=self.eps, out=u)
                if self.n_power_iterations > 0:
                    # See above on why we need to clone
                    u = u.clone()
                    v = v.clone()

        sigma = torch.dot(u, torch.mv(weight_mat, v))
        if self.rank_ratio > 0.9999:
            return weight / sigma
        weight_mat = weight_mat / sigma

        rank_1 = torch.ger(u, v)
        residual = weight_mat - rank_1
        frob = torch.norm(residual)
        num = math.sqrt(self.rank - 1)
        if frob > num:
            weight_mat = residual * num / frob + rank_1
        weight_ = self.reshape_matrix_to_weight(weight_mat, curr_shape)
        return weight_
//  [python]
def compute_weight(self, module, do_power_iteration):
        weight = getattr(module, self.name + '_orig')
        u = getattr(module, self.name + '_u')
        v = getattr(module, self.name + '_v')
        curr_shape = weight.shape
        weight_mat = self.reshape_weight_to_matrix(weight)

        if do_power_iteration:
            with torch.no_grad():
                for _ in range(self.n_power_iterations):
                    # Spectral norm of weight equals to `u^T W v`,
                    # where `u` and `v`
                    # are the first left and right singular vectors.
                    # This power iteration produces approximations
                    # of `u` and `v`.
                    v = normalize(torch.mv(weight_mat.t(), u),
                                  dim=0, eps=self.eps, out=v)
                    u = normalize(torch.mv(weight_mat, v),
                                  dim=0, eps=self.eps, out=u)
                if self.n_power_iterations > 0:
                    # See above on why we need to clone
                    u = u.clone()
                    v = v.clone()

        sigma = torch.dot(u, torch.mv(weight_mat, v))
        if self.rank_ratio > 0.9999:
            return weight / sigma
        weight_mat = weight_mat / sigma

        rank_1 = torch.ger(u, v)
        residual = weight_mat - rank_1
        frob = torch.linalg.norm(residual)
        num = math.sqrt(self.rank - 1)
        if frob > num:
            weight_mat = residual * num / frob + rank_1
        weight_ = self.reshape_matrix_to_weight(weight_mat, curr_shape)
        return weight_
//  [python]
def cost(X, w):
    wTX = torch.matmul(w.transpose(1,0), X)
    wwTX = torch.matmul(w, wTX)
    return torch.norm((X - wwTX)**2)
//  [python]
def cost(X, w):
    wTX = torch.matmul(w.transpose(1,0), X)
    wwTX = torch.matmul(w, wTX)
    return torch.linalg.norm((X - wwTX)**2)
//  [python]
def preprocessing(self,X):
        """
        Perform mean_subtraction using mean.vec-> apply transform.mat -> input length norm

        Parameters
        ----------
        X : Xvectors N X d

        Returns
        -------
        transformed xvectors N X D

        """
        
        xvecs = torch.transpose(X,1,2) # DX N
        dim = xvecs.shape[1]
        # preprocessing
        # mean subtraction
        xvecs = xvecs - self.mean_vec[np.newaxis,:,np.newaxis]
        # PCA transform
        xvecs = torch.bmm(self.transform_mat[np.newaxis],xvecs)
        l2_norm = torch.norm(xvecs, dim=1, keepdim=True)
        l2_norm = l2_norm/math.sqrt(dim)
        # l2_norm = np.linalg.norm(xvecs,axis=1,keepdims=True)
        xvecsnew = xvecs/l2_norm
        
        return torch.transpose(xvecsnew,1,2)
//  [python]
def preprocessing(self,X):
        """
        Perform mean_subtraction using mean.vec-> apply transform.mat -> input length norm

        Parameters
        ----------
        X : Xvectors N X d

        Returns
        -------
        transformed xvectors N X D

        """
        
        xvecs = torch.transpose(X,1,2) # DX N
        dim = xvecs.shape[1]
        # preprocessing
        # mean subtraction
        xvecs = xvecs - self.mean_vec[np.newaxis,:,np.newaxis]
        # PCA transform
        xvecs = torch.bmm(self.transform_mat[np.newaxis],xvecs)
        l2_norm = torch.linalg.norm(xvecs, dim=1, keepdim=True)
        l2_norm = l2_norm/math.sqrt(dim)
        # l2_norm = np.linalg.norm(xvecs,axis=1,keepdims=True)
        xvecsnew = xvecs/l2_norm
        
        return torch.transpose(xvecsnew,1,2)
//  [python]
def compute_affinity_matrix(self,X):
        """Compute the affinity matrix from data.

        Note that the range of affinity is [-1,1].

        Args:
            X: numpy array of shape (n_samples, n_features)

        Returns:
            affinity: numpy array of shape (n_samples, n_samples)
        """
        # Normalize the data.
        l2_norms = torch.norm(X, dim=2,keepdim=True)
        X_normalized = X / l2_norms
        # Compute cosine similarities. Range is [-1,1].

        cosine_similarities = torch.bmm(X_normalized, torch.transpose(X_normalized,1,2))
        # Compute the affinity. Range is [0,1].
        # Note that this step is not mentioned in the paper!
        # affinity = cosine_similarities
        affinity = cosine_similarities
        # affinity = (cosine_similarities + 1.0) / 2.0

        return affinity
//  [python]
def compute_affinity_matrix(self,X):
        """Compute the affinity matrix from data.

        Note that the range of affinity is [-1,1].

        Args:
            X: numpy array of shape (n_samples, n_features)

        Returns:
            affinity: numpy array of shape (n_samples, n_samples)
        """
        # Normalize the data.
        l2_norms = torch.linalg.norm(X, dim=2,keepdim=True)
        X_normalized = X / l2_norms
        # Compute cosine similarities. Range is [-1,1].

        cosine_similarities = torch.bmm(X_normalized, torch.transpose(X_normalized,1,2))
        # Compute the affinity. Range is [0,1].
        # Note that this step is not mentioned in the paper!
        # affinity = cosine_similarities
        affinity = cosine_similarities
        # affinity = (cosine_similarities + 1.0) / 2.0

        return affinity
//  [python]
def preprocessing(self,X):
        """
        Perform mean_subtraction using mean.vec-> apply transform.mat -> input length norm

        Parameters
        ----------
        X : Xvectors 1 X N X D

        Returns
        -------
        transformed xvectors 1 X N X D

        """
        
        # xvecs = torch.transpose(X,1,2) # 1 X DX N
        xvecs = X # 1X N X D
        dim = xvecs.shape[2]
        # preprocessing
        # mean subtraction
        xvecs = xvecs - self.mean_vec[np.newaxis,np.newaxis,:]
        # PCA transform        
        # xvecs = torch.bmm(self.transform_mat[np.newaxis],xvecs)
        # bp()
        xvecs = self.globalpca_out(xvecs)
        l2_norm = torch.norm(xvecs, dim=2, keepdim=True)
        l2_norm = l2_norm/math.sqrt(dim)
        # l2_norm = np.linalg.norm(xvecs,axis=1,keepdims=True)
        xvecsnew = xvecs/l2_norm
        
        return xvecsnew
//  [python]
def preprocessing(self,X):
        """
        Perform mean_subtraction using mean.vec-> apply transform.mat -> input length norm

        Parameters
        ----------
        X : Xvectors 1 X N X D

        Returns
        -------
        transformed xvectors 1 X N X D

        """
        
        # xvecs = torch.transpose(X,1,2) # 1 X DX N
        xvecs = X # 1X N X D
        dim = xvecs.shape[2]
        # preprocessing
        # mean subtraction
        xvecs = xvecs - self.mean_vec[np.newaxis,np.newaxis,:]
        # PCA transform        
        # xvecs = torch.bmm(self.transform_mat[np.newaxis],xvecs)
        # bp()
        xvecs = self.globalpca_out(xvecs)
        l2_norm = torch.linalg.norm(xvecs, dim=2, keepdim=True)
        l2_norm = l2_norm/math.sqrt(dim)
        # l2_norm = np.linalg.norm(xvecs,axis=1,keepdims=True)
        xvecsnew = xvecs/l2_norm
        
        return xvecsnew
//  [python]
def compute_affinity_matrix(self,X,attention=False):
        """Compute the affinity matrix from data.

        Note that the range of affinity is [0,1].

        Args:
            X: numpy array of shape (n_samples, n_features)

        Returns:
            affinity: numpy array of shape (n_samples, n_samples)
        """
        X = self.preprocessing(X) #output -N X D
        X = self.filepca_out(X)
        
        # X = X - torch.mean(X,dim=1,keepdim=True)        
        
        if type(attention)!=bool:            
            X = torch.bmm(attention,X)
            
        # Normalize the data.
        
        l2_norms = torch.norm(X, dim=2,keepdim=True)
        X_normalized = X / l2_norms
        # Compute cosine similarities. Range is [-1,1].

        cosine_similarities = torch.bmm(X_normalized, torch.transpose(X_normalized,1,2))
       
        affinity = cosine_similarities
        return affinity
//  [python]
def compute_affinity_matrix(self,X,attention=False):
        """Compute the affinity matrix from data.

        Note that the range of affinity is [0,1].

        Args:
            X: numpy array of shape (n_samples, n_features)

        Returns:
            affinity: numpy array of shape (n_samples, n_samples)
        """
        X = self.preprocessing(X) #output -N X D
        X = self.filepca_out(X)
        
        # X = X - torch.mean(X,dim=1,keepdim=True)        
        
        if type(attention)!=bool:            
            X = torch.bmm(attention,X)
            
        # Normalize the data.
        
        l2_norms = torch.linalg.norm(X, dim=2,keepdim=True)
        X_normalized = X / l2_norms
        # Compute cosine similarities. Range is [-1,1].

        cosine_similarities = torch.bmm(X_normalized, torch.transpose(X_normalized,1,2))
       
        affinity = cosine_similarities
        return affinity
//  [python]
def clip_grad(theta, max_norm):
        if max_norm < 0:
            return theta, torch.tensor(1.0, device=theta.device)
        total_norm = torch.norm(theta)
        clip_coef = max_norm / (total_norm + 1e-6)
        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
        theta.mul_(clip_coef_clamped)
        return theta, clip_coef_clamped
//  [python]
def clip_grad(theta, max_norm):
        if max_norm < 0:
            return theta, torch.tensor(1.0, device=theta.device)
        total_norm = torch.linalg.norm(theta)
        clip_coef = max_norm / (total_norm + 1e-6)
        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
        theta.mul_(clip_coef_clamped)
        return theta, clip_coef_clamped
//  [python]
def train(self):
        total_loss = 0
        total_tv_loss = 0
        total_orthogonal_constraint_loss = 0

        pbar = tqdm(range(args.max_iter_biased_norm_vec), dynamic_ncols=True)
        for idx in pbar:
            z = torch.rand(self.args.batch_size,
                           self.args.z_dim, device=self.device)

            z_projected = project_point_to_plane(z, self.w, self.b)
            z_traverse = traverse_along_normal_vec(
                z_projected, self.w, self.traverse_steps)

            out = self.model.decode(
                latent=z_traverse.reshape(-1, self.args.z_dim))
            prob = torch.sigmoid(self.biased_classifier(out)).reshape(
                self.args.batch_size, self.traverse_steps.shape[1])
            # total variation loss
            tv_loss = - \
                torch.log(1e-10 + torch.abs(prob[:, 1:] - prob[:, :-1]).mean())
            total_tv_loss += tv_loss.item()
            orthogonal_constraint_loss = (
                ((self.w / torch.norm(self.w)) @ self.unit_gt_other_normal_vectors) ** 2).mean()
            total_orthogonal_constraint_loss += orthogonal_constraint_loss.item()
            loss = tv_loss + self.args.lambda_orthogonal_constraint * orthogonal_constraint_loss

            self.optimizer.zero_grad()
            self.biased_classifier.zero_grad()
            self.model.model.zero_grad()
            loss.backward()
            self.optimizer.step()
            total_loss += loss.item()
            avg_loss = total_loss / (idx + 1)
            with torch.no_grad():
                cosine_sim = F.cosine_similarity(
                    self.w, self.gt_biased_normal_vec)
                target_cosine_sim = F.cosine_similarity(
                    self.w, self.gt_target_normal_vec)

            desc = 'loss: {:.3f}, normal_vec_cos: {:.3f}, ' \
                   'target_vec_cos_sim: {:.3f}'.format(avg_loss, cosine_sim.item(),
                                                       target_cosine_sim.item())
            pbar.set_description(desc)
//  [python]
def train(self):
        total_loss = 0
        total_tv_loss = 0
        total_orthogonal_constraint_loss = 0

        pbar = tqdm(range(args.max_iter_biased_norm_vec), dynamic_ncols=True)
        for idx in pbar:
            z = torch.rand(self.args.batch_size,
                           self.args.z_dim, device=self.device)

            z_projected = project_point_to_plane(z, self.w, self.b)
            z_traverse = traverse_along_normal_vec(
                z_projected, self.w, self.traverse_steps)

            out = self.model.decode(
                latent=z_traverse.reshape(-1, self.args.z_dim))
            prob = torch.sigmoid(self.biased_classifier(out)).reshape(
                self.args.batch_size, self.traverse_steps.shape[1])
            # total variation loss
            tv_loss = - \
                torch.log(1e-10 + torch.abs(prob[:, 1:] - prob[:, :-1]).mean())
            total_tv_loss += tv_loss.item()
            orthogonal_constraint_loss = (
                ((self.w / torch.linalg.norm(self.w)) @ self.unit_gt_other_normal_vectors) ** 2).mean()
            total_orthogonal_constraint_loss += orthogonal_constraint_loss.item()
            loss = tv_loss + self.args.lambda_orthogonal_constraint * orthogonal_constraint_loss

            self.optimizer.zero_grad()
            self.biased_classifier.zero_grad()
            self.model.model.zero_grad()
            loss.backward()
            self.optimizer.step()
            total_loss += loss.item()
            avg_loss = total_loss / (idx + 1)
            with torch.no_grad():
                cosine_sim = F.cosine_similarity(
                    self.w, self.gt_biased_normal_vec)
                target_cosine_sim = F.cosine_similarity(
                    self.w, self.gt_target_normal_vec)

            desc = 'loss: {:.3f}, normal_vec_cos: {:.3f}, ' \
                   'target_vec_cos_sim: {:.3f}'.format(avg_loss, cosine_sim.item(),
                                                       target_cosine_sim.item())
            pbar.set_description(desc)
//  [python]
def vanishing_point(X, vp):
    lines = X[..., 6:9]
    centroids = X[..., 9:12]

    constrained_lines = torch.cross(centroids, vp.unsqueeze(-2), dim=-1)
    con_line_norms = torch.norm(constrained_lines[..., 0:2], dim=-1, keepdim=True)
    constrained_lines = constrained_lines / (con_line_norms + 1e-8)
    line_norms = torch.norm(lines[..., 0:2], dim=-1, keepdim=True)
    lines = lines / (line_norms + 1e-8)
    distances = 1 - torch.abs((lines[..., 0:2] * constrained_lines[..., 0:2]).sum(dim=-1))

    return distances
//  [python]
def vanishing_point(X, vp):
    lines = X[..., 6:9]
    centroids = X[..., 9:12]

    constrained_lines = torch.cross(centroids, vp.unsqueeze(-2), dim=-1)
    con_line_norms = torch.linalg.norm(constrained_lines[..., 0:2], dim=-1, keepdim=True)
    constrained_lines = constrained_lines / (con_line_norms + 1e-8)
    line_norms = torch.linalg.norm(lines[..., 0:2], dim=-1, keepdim=True)
    lines = lines / (line_norms + 1e-8)
    distances = 1 - torch.abs((lines[..., 0:2] * constrained_lines[..., 0:2]).sum(dim=-1))

    return distances
//  [python]
def loss_laplacian(deformed_vert, face_3d, temp_L):
    n_batch = deformed_vert.shape[0]
    curve_gt = torch.norm(temp_L.view(-1, 3), p=2, dim=1).float()
    L_norm = get_laplacian(deformed_vert, face_3d)[0]
    Lx = L_norm.mm(deformed_vert.view(-1, 3))
    repeated_gt = curve_gt.repeat(n_batch)
    loss = ((torch.norm(Lx, p=2, dim=1).float() - repeated_gt)**2).mean()
    return loss
//  [python]
def loss_laplacian(deformed_vert, face_3d, temp_L):
    n_batch = deformed_vert.shape[0]
    curve_gt = torch.linalg.norm(temp_L.view(-1, 3), ord=2, dim=1).float()
    L_norm = get_laplacian(deformed_vert, face_3d)[0]
    Lx = L_norm.mm(deformed_vert.view(-1, 3))
    repeated_gt = curve_gt.repeat(n_batch)
    loss = ((torch.linalg.norm(Lx, ord=2, dim=1).float() - repeated_gt)**2).mean()
    return loss
//  [python]
def sampson_distance(X, f, squared=False):
    f_dims = f.size()
    batch_dims = list(f_dims)[:-1]
    F_dims = tuple(list(f_dims)[:-1] + [3, 3])

    degenerate = torch.norm(f, dim=-1) < 1e-8

    N = X.size(-2)

    F = f.view(F_dims)[..., None, :, :]

    x1 = X[..., 0:2]
    x2 = X[..., 2:4]
    X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X1[..., 0:2, 0] = x1
    X2[..., 0:2, 0] = x2

    Fx1 = F @ X1
    Fx2 = F.transpose(-1, -2) @ X2

    xFx = X2 * Fx1
    xFx = torch.sum(xFx[..., 0], dim=-1) ** 2

    denom = Fx1[..., 0, 0] ** 2 + Fx1[..., 1, 0] ** 2 + Fx2[..., 0, 0] ** 2 + Fx2[..., 1, 0] ** 2
    denom = torch.clamp(denom, min=1e-8)

    sq_distances = xFx / denom

    sq_distances = sq_distances + degenerate[..., None].float() * 1e6

    if squared:
        return sq_distances
    else:
        return torch.sqrt(sq_distances)
//  [python]
def sampson_distance(X, f, squared=False):
    f_dims = f.size()
    batch_dims = list(f_dims)[:-1]
    F_dims = tuple(list(f_dims)[:-1] + [3, 3])

    degenerate = torch.linalg.norm(f, dim=-1) < 1e-8

    N = X.size(-2)

    F = f.view(F_dims)[..., None, :, :]

    x1 = X[..., 0:2]
    x2 = X[..., 2:4]
    X1 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X2 = torch.ones(batch_dims + [N, 3, 1], device=X.device)
    X1[..., 0:2, 0] = x1
    X2[..., 0:2, 0] = x2

    Fx1 = F @ X1
    Fx2 = F.transpose(-1, -2) @ X2

    xFx = X2 * Fx1
    xFx = torch.sum(xFx[..., 0], dim=-1) ** 2

    denom = Fx1[..., 0, 0] ** 2 + Fx1[..., 1, 0] ** 2 + Fx2[..., 0, 0] ** 2 + Fx2[..., 1, 0] ** 2
    denom = torch.clamp(denom, min=1e-8)

    sq_distances = xFx / denom

    sq_distances = sq_distances + degenerate[..., None].float() * 1e6

    if squared:
        return sq_distances
    else:
        return torch.sqrt(sq_distances)
//  [python]
def get_cam_rays(H, W, K):
    grid = create_meshgrid(H, W, normalized_coordinates=False)[0]
    i, j = grid.unbind(-1)
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    directions = directions / torch.norm(directions, dim=-1, keepdim=True)
    return directions.numpy()
//  [python]
def get_cam_rays(H, W, K):
    grid = create_meshgrid(H, W, normalized_coordinates=False)[0]
    i, j = grid.unbind(-1)
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    directions = directions / torch.linalg.norm(directions, dim=-1, keepdim=True)
    return directions.numpy()
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor,
                                 num_epochs):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples - x_equilibrium) *
                          (lqr_S @ (x_samples - x_equilibrium).T).T,
                          dim=1).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.norm(
            R @ (x - x_equilibrium).T, p=1, dim=0).reshape((-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=num_epochs,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor,
                                 num_epochs):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples - x_equilibrium) *
                          (lqr_S @ (x_samples - x_equilibrium).T).T,
                          dim=1).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.linalg.norm(
            R @ (x - x_equilibrium).T, ord=1, dim=0).reshape((-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=num_epochs,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def patch(self, model, empty_conditioning, neg_scale):
        m = model.clone()
        nocond = ldm_patched.modules.sample.convert_cond(empty_conditioning)

        def cfg_function(args):
            model = args["model"]
            noise_pred_pos = args["cond_denoised"]
            noise_pred_neg = args["uncond_denoised"]
            cond_scale = args["cond_scale"]
            x = args["input"]
            sigma = args["sigma"]
            model_options = args["model_options"]
            nocond_processed = ldm_patched.modules.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative")

            (noise_pred_nocond, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options)

            pos = noise_pred_pos - noise_pred_nocond
            neg = noise_pred_neg - noise_pred_nocond
            perp = ((torch.mul(pos, neg).sum())/(torch.norm(neg)**2)) * neg
            perp_neg = perp * neg_scale
            cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
            cfg_result = x - cfg_result
            return cfg_result

        m.set_model_sampler_cfg_function(cfg_function)

        return (m, )
//  [python]
def patch(self, model, empty_conditioning, neg_scale):
        m = model.clone()
        nocond = ldm_patched.modules.sample.convert_cond(empty_conditioning)

        def cfg_function(args):
            model = args["model"]
            noise_pred_pos = args["cond_denoised"]
            noise_pred_neg = args["uncond_denoised"]
            cond_scale = args["cond_scale"]
            x = args["input"]
            sigma = args["sigma"]
            model_options = args["model_options"]
            nocond_processed = ldm_patched.modules.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative")

            (noise_pred_nocond, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options)

            pos = noise_pred_pos - noise_pred_nocond
            neg = noise_pred_neg - noise_pred_nocond
            perp = ((torch.mul(pos, neg).sum())/(torch.linalg.norm(neg)**2)) * neg
            perp_neg = perp * neg_scale
            cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
            cfg_result = x - cfg_result
            return cfg_result

        m.set_model_sampler_cfg_function(cfg_function)

        return (m, )
//  [python]
def angular_loss(predicted, gt, shrink=True):
  """Computes angular error between predicted and gt illuminant color(s)."""
  cossim = torch.clamp(torch.sum(predicted * gt, dim=1) / (
      torch.norm(predicted, dim=1) * torch.norm(gt, dim=1) + EPS), -1, 1.)
  if shrink:
    angle = torch.acos(cossim * 0.999999)
  else:
    angle = torch.acos(cossim)
  a_error = 180.0 / PI * angle
  a_error = torch.sum(a_error) / a_error.shape[0]
  return a_error
//  [python]
def angular_loss(predicted, gt, shrink=True):
  """Computes angular error between predicted and gt illuminant color(s)."""
  cossim = torch.clamp(torch.sum(predicted * gt, dim=1) / (
      torch.linalg.norm(predicted, dim=1) * torch.linalg.norm(gt, dim=1) + EPS), -1, 1.)
  if shrink:
    angle = torch.acos(cossim * 0.999999)
  else:
    angle = torch.acos(cossim)
  a_error = 180.0 / PI * angle
  a_error = torch.sum(a_error) / a_error.shape[0]
  return a_error
//  [python]
def test_dense_solvers(damp_as_tensor):
    void_objective = th.Objective()
    void_ordering = th.VariableOrdering(void_objective, default_order=False)
    for solver_cls in [th.LUDenseSolver, th.CholeskyDenseSolver]:
        solver = solver_cls(
            void_objective, linearization_kwargs={"ordering": void_ordering}
        )
        (
            solver.linearization._AtA,
            solver.linearization._Atb,
            x,
        ) = _create_linear_system()
        solved_x = solver.solve().squeeze()
        error = torch.norm(x - solved_x, p=float("inf"))
        assert error < 1e-4

        # Test damping
        damping = (
            0.2
            if not damp_as_tensor
            else 0.1 * torch.arange(x.shape[0], device=x.device, dtype=x.dtype)
        )
        damping_eps = 1e-3
        solved_x_with_ellipsoidal_damp = solver.solve(
            damping=damping, ellipsoidal_damping=True, damping_eps=damping_eps
        ).squeeze()
        solved_x_with_spherical_damp = solver.solve(
            damping=damping, ellipsoidal_damping=False
        ).squeeze()
        AtA_ellipsoidal_damp = solver.linearization.AtA.clone()
        AtA_spherical_damp = solver.linearization.AtA.clone()
        batch_size, n, _ = AtA_ellipsoidal_damp.shape
        # Add damping (using loop to make it more visually explicit)
        for i in range(batch_size):
            for j in range(n):
                damp = damping if isinstance(damping, float) else damping[i]
                AtA_ellipsoidal_damp[i, j, j] *= 1 + damp
                AtA_ellipsoidal_damp[i, j, j] += damping_eps
                AtA_spherical_damp[i, j, j] += damp
        # AtA_ellipsoidal_damp += damping_eps
        Atb_check_ellipsoidal = AtA_ellipsoidal_damp.bmm(
            solved_x_with_ellipsoidal_damp.unsqueeze(-1)
        )
        Atb_check_spherical = AtA_spherical_damp.bmm(
            solved_x_with_spherical_damp.unsqueeze(-1)
        )
        error = torch.norm(
            solver.linearization.Atb - Atb_check_ellipsoidal, p=float("inf")
        )
        assert error < 1e-4
        error = torch.norm(
            solver.linearization.Atb - Atb_check_spherical, p=float("inf")
        )
        assert error < 1e-4
//  [python]
def test_dense_solvers(damp_as_tensor):
    void_objective = th.Objective()
    void_ordering = th.VariableOrdering(void_objective, default_order=False)
    for solver_cls in [th.LUDenseSolver, th.CholeskyDenseSolver]:
        solver = solver_cls(
            void_objective, linearization_kwargs={"ordering": void_ordering}
        )
        (
            solver.linearization._AtA,
            solver.linearization._Atb,
            x,
        ) = _create_linear_system()
        solved_x = solver.solve().squeeze()
        error = torch.linalg.norm(x - solved_x, p=float("inf"))
        assert error < 1e-4

        # Test damping
        damping = (
            0.2
            if not damp_as_tensor
            else 0.1 * torch.arange(x.shape[0], device=x.device, dtype=x.dtype)
        )
        damping_eps = 1e-3
        solved_x_with_ellipsoidal_damp = solver.solve(
            damping=damping, ellipsoidal_damping=True, damping_eps=damping_eps
        ).squeeze()
        solved_x_with_spherical_damp = solver.solve(
            damping=damping, ellipsoidal_damping=False
        ).squeeze()
        AtA_ellipsoidal_damp = solver.linearization.AtA.clone()
        AtA_spherical_damp = solver.linearization.AtA.clone()
        batch_size, n, _ = AtA_ellipsoidal_damp.shape
        # Add damping (using loop to make it more visually explicit)
        for i in range(batch_size):
            for j in range(n):
                damp = damping if isinstance(damping, float) else damping[i]
                AtA_ellipsoidal_damp[i, j, j] *= 1 + damp
                AtA_ellipsoidal_damp[i, j, j] += damping_eps
                AtA_spherical_damp[i, j, j] += damp
        # AtA_ellipsoidal_damp += damping_eps
        Atb_check_ellipsoidal = AtA_ellipsoidal_damp.bmm(
            solved_x_with_ellipsoidal_damp.unsqueeze(-1)
        )
        Atb_check_spherical = AtA_spherical_damp.bmm(
            solved_x_with_spherical_damp.unsqueeze(-1)
        )
        error = torch.linalg.norm(
            solver.linearization.Atb - Atb_check_ellipsoidal, p=float("inf")
        )
        assert error < 1e-4
        error = torch.linalg.norm(
            solver.linearization.Atb - Atb_check_spherical, p=float("inf")
        )
        assert error < 1e-4
//  [python]
def forward(self, mel):
        mels = mel.unfold(1, self.hp['window'], self.hp['stride'])
        mels = mels.permute(1, 2, 0)
        x, _ = self.lstm(mels)
        x = x[:, -1, :]
        x = self.proj(x)
        x = x / torch.norm(x, p=2, dim=1, keepdim=True)
        x = x.mean(dim=0)
        if x.norm(p=2) != 0:
            x = x / x.norm(p=2)
        return x
//  [python]
def forward(self, mel):
        mels = mel.unfold(1, self.hp['window'], self.hp['stride'])
        mels = mels.permute(1, 2, 0)
        x, _ = self.lstm(mels)
        x = x[:, -1, :]
        x = self.proj(x)
        x = x / torch.linalg.norm(x, p=2, dim=1, keepdim=True)
        x = x.mean(dim=0)
        if torch.linalg.norm(x, p=2) != 0:
            x = x / torch.linalg.norm(x, p=2)
        return x
//  [python]
def forward(self, mels: torch.FloatTensor):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        :param mels: a batch of mel spectrograms of same duration as a float32 tensor of shape 
        (batch_size, n_frames, n_channels) 
        :return: the embeddings as a float 32 tensor of shape (batch_size, embedding_size). 
        Embeddings are positive and L2-normed, thus they lay in the range [0, 1].
        """
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def forward(self, mels: torch.FloatTensor):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        :param mels: a batch of mel spectrograms of same duration as a float32 tensor of shape 
        (batch_size, n_frames, n_channels) 
        :return: the embeddings as a float 32 tensor of shape (batch_size, embedding_size). 
        Embeddings are positive and L2-normed, thus they lay in the range [0, 1].
        """
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def solve(self, qobjs, kobjs, f0):
        r"""Regularized Hough matching"""
        # Unpack hyperpixels
        m0 = qobjs.mask.float()
        f0 = f0.float()
        f1 = kobjs['feature'].to(m0).float()
        m1 = kobjs['mask'].to(m0).float()
       
        fg_mask = m0.reshape(m0.shape[0], -1, 1) * m1.reshape(m1.shape[0], 1, -1)
        bg_mask = (1 - m0).reshape(m0.shape[0], -1, 1) * (1 - m1).reshape(m1.shape[0], 1, -1)
        
        m0 = F.interpolate(m0.unsqueeze(1), (f0.shape[2], f0.shape[3]), mode='bilinear', align_corners=False).squeeze(1)
        m1 = F.interpolate(m1.unsqueeze(1), (f1.shape[2], f1.shape[3]), mode='bilinear', align_corners=False).squeeze(1)
        shape = f0.shape[2], f0.shape[3]

        m0 = m0.reshape(m0.shape[0], -1)
        m1 = m1.reshape(m1.shape[0], -1)
        f0 = f0.reshape(f0.shape[0], f0.shape[1], -1).transpose(2, 1)
        f1 = f1.reshape(f1.shape[0], f1.shape[1], -1)

        f0_norm = torch.norm(f0, p=2, dim=2, keepdim=True) + 1e-4
        f1_norm = torch.norm(f1, p=2, dim=1, keepdim=True) + 1e-4
        with autocast(enabled=False):
            Cu = torch.matmul((f0 / f0_norm), (f1 / f1_norm))

        eye = torch.eye(shape[0] * shape[1]).to(f0).reshape(1, -1, shape[0], shape[1])
        dist_mask = F.max_pool2d(eye, kernel_size=self.dist_kernel, stride=1, padding=self.dist_kernel//2).reshape(1, shape[0] * shape[1],
                                                                                    shape[0] * shape[1]).transpose(2, 1)
        with torch.no_grad():
            C = Cu.clone() * dist_mask

        for i in range(self.num_iter):
            pairwise_votes = C.clone()
            for _ in range(self.num_smooth_iter):
                pairwise_votes = self.pass_message(pairwise_votes, (shape[0], shape[1]))
                pairwise_votes = pairwise_votes / (pairwise_votes.sum(2, keepdim=True) + 1e-4)

            max_val, _ = pairwise_votes.max(2, keepdim=True)

            C = Cu + pairwise_votes
            C = C / (C.sum(2, keepdim=True) + 1e-4)

        return Cu, C, fg_mask, bg_mask
//  [python]
def solve(self, qobjs, kobjs, f0):
        r"""Regularized Hough matching"""
        # Unpack hyperpixels
        m0 = qobjs.mask.float()
        f0 = f0.float()
        f1 = kobjs['feature'].to(m0).float()
        m1 = kobjs['mask'].to(m0).float()
       
        fg_mask = m0.reshape(m0.shape[0], -1, 1) * m1.reshape(m1.shape[0], 1, -1)
        bg_mask = (1 - m0).reshape(m0.shape[0], -1, 1) * (1 - m1).reshape(m1.shape[0], 1, -1)
        
        m0 = F.interpolate(m0.unsqueeze(1), (f0.shape[2], f0.shape[3]), mode='bilinear', align_corners=False).squeeze(1)
        m1 = F.interpolate(m1.unsqueeze(1), (f1.shape[2], f1.shape[3]), mode='bilinear', align_corners=False).squeeze(1)
        shape = f0.shape[2], f0.shape[3]

        m0 = m0.reshape(m0.shape[0], -1)
        m1 = m1.reshape(m1.shape[0], -1)
        f0 = f0.reshape(f0.shape[0], f0.shape[1], -1).transpose(2, 1)
        f1 = f1.reshape(f1.shape[0], f1.shape[1], -1)

        f0_norm = torch.linalg.norm(f0, ord=2, dim=2, keepdim=True) + 1e-4
        f1_norm = torch.linalg.norm(f1, ord=2, dim=1, keepdim=True) + 1e-4
        with autocast(enabled=False):
            Cu = torch.matmul((f0 / f0_norm), (f1 / f1_norm))

        eye = torch.eye(shape[0] * shape[1]).to(f0).reshape(1, -1, shape[0], shape[1])
        dist_mask = F.max_pool2d(eye, kernel_size=self.dist_kernel, stride=1, padding=self.dist_kernel//2).reshape(1, shape[0] * shape[1],
                                                                                    shape[0] * shape[1]).transpose(2, 1)
        with torch.no_grad():
            C = Cu.clone() * dist_mask

        for i in range(self.num_iter):
            pairwise_votes = C.clone()
            for _ in range(self.num_smooth_iter):
                pairwise_votes = self.pass_message(pairwise_votes, (shape[0], shape[1]))
                pairwise_votes = pairwise_votes / (pairwise_votes.sum(2, keepdim=True) + 1e-4)

            max_val, _ = pairwise_votes.max(2, keepdim=True)

            C = Cu + pairwise_votes
            C = C / (C.sum(2, keepdim=True) + 1e-4)

        return Cu, C, fg_mask, bg_mask
//  [python]
def training_step(self, batch, batch_idx):
        batch.pos = batch.positions

        emb_2d = batch.emb_2d
        # _, emb_3d = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True)
        _, emb_3d, emb_3d_local = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True, True)
        
        loss_cl = self.loss_cl(emb_3d, emb_3d_local, batch.n_nodes)

        # spatial graphs
        batch.edge_index = batch.radius_edge_index
        edge_attr = torch.exp( - torch.norm(batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]], dim=1) )
        batch.edge_attr = torch.einsum('i,j->ij', edge_attr, torch.linspace(1, 5, 16).to(edge_attr.device))
        batch.n_edges = batch.n_edges_radius

        atom_pred_list, extra_output = self.decoder_2dto3d(batch, emb_3d)

        gt_poss = tgeom.utils.unbatch(batch.pos, batch.batch)
        atom_pred_list = tgeom.utils.unbatch(atom_pred_list[-1], batch.batch)

        loss = 0
        for pred_pos, gt_pos in zip(atom_pred_list, gt_poss):
            pred_pos = self.kbash_alignment(pred_pos, gt_pos)
            _loss = (pred_pos - gt_pos).pow(2).mean()
            loss += _loss

            # print(_loss)
            # pdb.set_trace()
        loss /= len(atom_pred_list)

        # avoiding collapse
        emb_3d_norm = emb_3d.norm(dim=1)
        sim_matrix = torch.einsum('ik,jk->ij', emb_3d, emb_3d) / torch.einsum('i,j->ij', emb_3d_norm, emb_3d_norm)
        sim_matrix = torch.exp(sim_matrix)
        pos_sim = sim_matrix[range(emb_3d.shape[0]), range(emb_3d.shape[0])]
        loss_neg_contrast = torch.log( (sim_matrix.sum(dim=1) - pos_sim) ).mean()
        loss += loss_neg_contrast * 0.1 + + 1e-3 * loss_cl

        # loss, loss_dict = self.decoder_2dto3d.compute_loss(atom_pred_list, extra_output, batch)
        self.log('train_loss', loss, batch_size=batch.num_graphs, on_epoch=True, sync_dist=True)

        return loss
//  [python]
def training_step(self, batch, batch_idx):
        batch.pos = batch.positions

        emb_2d = batch.emb_2d
        # _, emb_3d = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True)
        _, emb_3d, emb_3d_local = self.encoder_3d(batch.x, batch.pos, batch.radius_edge_index, batch.batch, True, True)
        
        loss_cl = self.loss_cl(emb_3d, emb_3d_local, batch.n_nodes)

        # spatial graphs
        batch.edge_index = batch.radius_edge_index
        edge_attr = torch.exp( - torch.linalg.norm(batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]], dim=1) )
        batch.edge_attr = torch.einsum('i,j->ij', edge_attr, torch.linspace(1, 5, 16).to(edge_attr.device))
        batch.n_edges = batch.n_edges_radius

        atom_pred_list, extra_output = self.decoder_2dto3d(batch, emb_3d)

        gt_poss = tgeom.utils.unbatch(batch.pos, batch.batch)
        atom_pred_list = tgeom.utils.unbatch(atom_pred_list[-1], batch.batch)

        loss = 0
        for pred_pos, gt_pos in zip(atom_pred_list, gt_poss):
            pred_pos = self.kbash_alignment(pred_pos, gt_pos)
            _loss = (pred_pos - gt_pos).pow(2).mean()
            loss += _loss

            # print(_loss)
            # pdb.set_trace()
        loss /= len(atom_pred_list)

        # avoiding collapse
        emb_3d_norm = emb_3d.norm(dim=1)
        sim_matrix = torch.einsum('ik,jk->ij', emb_3d, emb_3d) / torch.einsum('i,j->ij', emb_3d_norm, emb_3d_norm)
        sim_matrix = torch.exp(sim_matrix)
        pos_sim = sim_matrix[range(emb_3d.shape[0]), range(emb_3d.shape[0])]
        loss_neg_contrast = torch.log( (sim_matrix.sum(dim=1) - pos_sim) ).mean()
        loss += loss_neg_contrast * 0.1 + + 1e-3 * loss_cl

        # loss, loss_dict = self.decoder_2dto3d.compute_loss(atom_pred_list, extra_output, batch)
        self.log('train_loss', loss, batch_size=batch.num_graphs, on_epoch=True, sync_dist=True)

        return loss
//  [python]
def get_mel(self, y):
        if self.pad and y.shape[-1] < 14000:
            y = F.pad(y, (0, 14000 - y.shape[-1]))
        window = torch.hann_window(self.win_length)
        y = torch.stft(y, n_fft=self.embedder_params['n_fft'], hop_length=self.hop_length, win_length=self.win_length, window=window)
        magnitudes = torch.norm(y, dim=-1, p=2) ** 2
        mel = torch.log10(self.mel_basis @ magnitudes + 1e-06)
        return mel
//  [python]
def get_mel(self, y):
        if self.pad and y.shape[-1] < 14000:
            y = F.pad(y, (0, 14000 - y.shape[-1]))
        window = torch.hann_window(self.win_length)
        y = torch.stft(y, n_fft=self.embedder_params['n_fft'], hop_length=self.hop_length, win_length=self.win_length, window=window)
        magnitudes = torch.linalg.norm(y, dim=-1, ord=2) ** 2
        mel = torch.log10(self.mel_basis @ magnitudes + 1e-06)
        return mel
//  [python]
def forward(self, inputs):
        dvecs = []
        for wav in inputs:
            mel = self.get_mel(wav)
            if mel.dim() == 3:
                mel = mel.squeeze(0)
            dvecs += [self.embedder(mel)]
        dvecs = torch.stack(dvecs)
        dvec = torch.mean(dvecs, dim=0)
        dvec = dvec / torch.norm(dvec)
        return dvec
//  [python]
def forward(self, inputs):
        dvecs = []
        for wav in inputs:
            mel = self.get_mel(wav)
            if mel.dim() == 3:
                mel = mel.squeeze(0)
            dvecs += [self.embedder(mel)]
        dvecs = torch.stack(dvecs)
        dvec = torch.mean(dvecs, dim=0)
        dvec = dvec / torch.linalg.norm(dvec)
        return dvec
//  [python]
def forward(self, x):

        def norm(t):
            n = torch.norm(t, dim=-1, keepdim=True).clamp(min=self.eps)
            return t / n * self.g
        return map_first_tuple_or_el(x, norm)
//  [python]
def forward(self, x):

        def norm(t):
            n = torch.linalg.norm(t, dim=-1, keepdim=True).clamp(min=self.eps)
            return t / n * self.g
        return map_first_tuple_or_el(x, norm)
//  [python]
def update(self, minibatches, opt, sch):
        all_x = torch.cat([data[0].cuda().float() for data in minibatches])
        all_y = torch.cat([data[1].cuda().long() for data in minibatches])
        with torch.no_grad():
            all_x1 = torch.angle(torch.fft.fftn(all_x, dim=(2, 3)))
            tfea = self.teab(self.teaf(all_x1)).detach()

        all_z = self.bottleneck(self.featurizer(all_x))
        loss1 = F.cross_entropy(self.classifier(all_z), all_y)

        loss2 = F.mse_loss(all_z[:, :self.tfbd], tfea)*self.args.alpha
        if self.args.disttype == '2-norm':
            loss3 = -F.mse_loss(all_z[:, :self.tfbd],
                                all_z[:, self.tfbd:])*self.args.beta
        elif self.args.disttype == 'norm-2-norm':
            loss3 = -F.mse_loss(all_z[:, :self.tfbd]/torch.norm(all_z[:, :self.tfbd], dim=1, keepdim=True),
                                all_z[:, self.tfbd:]/torch.norm(all_z[:, self.tfbd:], dim=1, keepdim=True))*self.args.beta
        elif self.args.disttype == 'norm-1-norm':
            loss3 = -F.l1_loss(all_z[:, :self.tfbd]/torch.norm(all_z[:, :self.tfbd], dim=1, keepdim=True),
                               all_z[:, self.tfbd:]/torch.norm(all_z[:, self.tfbd:], dim=1, keepdim=True))*self.args.beta
        elif self.args.disttype == 'cos':
            loss3 = torch.mean(F.cosine_similarity(
                all_z[:, :self.tfbd], all_z[:, self.tfbd:]))*self.args.beta
        loss4 = 0
        if len(minibatches) > 1:
            for i in range(len(minibatches)-1):
                for j in range(i+1, len(minibatches)):
                    loss4 += self.coral(all_z[i*self.args.batch_size:(i+1)*self.args.batch_size, self.tfbd:],
                                        all_z[j*self.args.batch_size:(j+1)*self.args.batch_size, self.tfbd:])
            loss4 = loss4*2/(len(minibatches) *
                             (len(minibatches)-1))*self.args.lam
        else:
            loss4 = self.coral(all_z[:self.args.batch_size//2, self.tfbd:],
                               all_z[self.args.batch_size//2:, self.tfbd:])
            loss4 = loss4*self.args.lam

        loss = loss1+loss2+loss3+loss4
        opt.zero_grad()
        loss.backward()
        opt.step()
        if sch:
            sch.step()
        return {'class': loss1.item(), 'dist': (loss2).item(), 'exp': (loss3).item(), 'align': loss4.item(), 'total': loss.item()}