jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, x):
        if self.norm:
            weights = self.context_vector(x) / torch.linalg.norm(self.context_vector.weight, p=2)
            weights = torch.add(torch.div(weights, 2.), .5)
        else:
            x_tr = self.activation(self.transform(x))
            weights = self.context_vector(x_tr)
            weights = torch.sigmoid(weights)
        x = x * weights
        return x, weights
//  [python]
def CPC_ALS(T, mask, R, iteration, device):
    """
    CPC-ALS:
        call <maskOptimizer> for the sub-iteration
    """

    # preparation
    loss_list, rec_list = [], []
    d1, d2, d3 = T.shape

    A1 = torch.FloatTensor(np.random.random((d1, R))).to(device)
    A2 = torch.FloatTensor(np.random.random((d2, R))).to(device)
    A3 = torch.FloatTensor(np.random.random((d3, R))).to(device)

    T_ = mask * T

    # ALS
    tic = time.time()
    for i in range(5):

        # sub-iteration
        A1 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,jr,kr->ir',T_,A2,A3), 0, 1e-5, device)
        A2 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,ir,kr->jr',T_,A1,A3), 1, 1e-5, device)
        A3 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,ir,jr->kr',T_,A1,A2), 2, 1e-5, device)

        # loss 
        rec = torch.einsum('ir,jr,kr->ijk',A1,A2,A3)
        LOSS = torch.norm(mask * (rec - T))/ torch.norm(mask * T)
        recLOSS = 1 - torch.norm(rec - T) / torch.norm(T)
        print ('{}/{}'.format(i, iteration), 'loss:', LOSS.item(), 'rec:', recLOSS.item(), 'time span:', time.time() - tic)
        
        # collect loss
        tic = time.time() 
        loss_list.append(LOSS)
        rec_list.append(recLOSS)

    return A1, A2, A3, loss_list, rec_list
//  [python]
def CPC_ALS(T, mask, R, iteration, device):
    """
    CPC-ALS:
        call <maskOptimizer> for the sub-iteration
    """

    # preparation
    loss_list, rec_list = [], []
    d1, d2, d3 = T.shape

    A1 = torch.FloatTensor(np.random.random((d1, R))).to(device)
    A2 = torch.FloatTensor(np.random.random((d2, R))).to(device)
    A3 = torch.FloatTensor(np.random.random((d3, R))).to(device)

    T_ = mask * T

    # ALS
    tic = time.time()
    for i in range(5):

        # sub-iteration
        A1 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,jr,kr->ir',T_,A2,A3), 0, 1e-5, device)
        A2 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,ir,kr->jr',T_,A1,A3), 1, 1e-5, device)
        A3 = maskOptimizer(mask, [A1, A2, A3], torch.einsum('ijk,ir,jr->kr',T_,A1,A2), 2, 1e-5, device)

        # loss 
        rec = torch.einsum('ir,jr,kr->ijk',A1,A2,A3)
        LOSS = torch.linalg.norm(mask * (rec - T))/ torch.linalg.norm(mask * T)
        recLOSS = 1 - torch.linalg.norm(rec - T) / torch.linalg.norm(T)
        print ('{}/{}'.format(i, iteration), 'loss:', LOSS.item(), 'rec:', recLOSS.item(), 'time span:', time.time() - tic)
        
        # collect loss
        tic = time.time() 
        loss_list.append(LOSS)
        rec_list.append(recLOSS)

    return A1, A2, A3, loss_list, rec_list
//  [python]
def precompute_terms(cls, base_lazy_tensor, eager_rhs, logdet_terms=True, include_tmats=True):
        """
        Computes the solves, probe vectors, probe_vector norms, probe vector solves, and probe vector
        tridiagonal matrices to construct a CachedCGLazyTensor

        Set logdet_terms to False if you are not going to compute the logdet of the LazyTensor
        """
        with torch.no_grad():
            if logdet_terms:
                # Generate probe vectors
                num_random_probes = settings.num_trace_samples.value()
                probe_vectors = torch.empty(
                    base_lazy_tensor.matrix_shape[-1],
                    num_random_probes,
                    dtype=base_lazy_tensor.dtype,
                    device=base_lazy_tensor.device,
                )
                probe_vectors.bernoulli_().mul_(2).add_(-1)
                probe_vectors = probe_vectors.expand(
                    *base_lazy_tensor.batch_shape, base_lazy_tensor.matrix_shape[-1], num_random_probes
                )
                probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
                probe_vectors = probe_vectors.div(probe_vector_norms)

                # Compute solves
                if include_tmats:
                    all_solves, probe_vector_tmats, = base_lazy_tensor._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_lazy_tensor._preconditioner()[0],
                        num_tridiag=probe_vectors.size(-1),
                    )
                else:
                    all_solves = base_lazy_tensor._solve(
                        torch.cat([probe_vectors, eager_rhs], -1), preconditioner=base_lazy_tensor._preconditioner()[0]
                    )
                    probe_vector_tmats = torch.tensor([])
                probe_vector_solves = all_solves[..., : probe_vectors.size(-1)].detach()
                solves = all_solves[..., probe_vectors.size(-1) :]

                return (
                    solves.detach(),
                    probe_vectors.detach(),
                    probe_vector_norms.detach(),
                    probe_vector_solves.detach(),
                    probe_vector_tmats.detach(),
                )

            else:
                # Compute solves
                if settings.fast_computations.log_prob.on():
                    solves = base_lazy_tensor._solve(eager_rhs, preconditioner=base_lazy_tensor._preconditioner()[0])
                else:
                    solves = base_lazy_tensor.cholesky()._cholesky_solve(eager_rhs)
                dtype = solves.dtype
                device = solves.device
                return (
                    solves.detach(),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                )
//  [python]
def precompute_terms(cls, base_lazy_tensor, eager_rhs, logdet_terms=True, include_tmats=True):
        """
        Computes the solves, probe vectors, probe_vector norms, probe vector solves, and probe vector
        tridiagonal matrices to construct a CachedCGLazyTensor

        Set logdet_terms to False if you are not going to compute the logdet of the LazyTensor
        """
        with torch.no_grad():
            if logdet_terms:
                # Generate probe vectors
                num_random_probes = settings.num_trace_samples.value()
                probe_vectors = torch.empty(
                    base_lazy_tensor.matrix_shape[-1],
                    num_random_probes,
                    dtype=base_lazy_tensor.dtype,
                    device=base_lazy_tensor.device,
                )
                probe_vectors.bernoulli_().mul_(2).add_(-1)
                probe_vectors = probe_vectors.expand(
                    *base_lazy_tensor.batch_shape, base_lazy_tensor.matrix_shape[-1], num_random_probes
                )
                probe_vector_norms = torch.linalg.norm(probe_vectors, 2, dim=-2, keepdim=True)
                probe_vectors = probe_vectors.div(probe_vector_norms)

                # Compute solves
                if include_tmats:
                    all_solves, probe_vector_tmats, = base_lazy_tensor._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_lazy_tensor._preconditioner()[0],
                        num_tridiag=probe_vectors.size(-1),
                    )
                else:
                    all_solves = base_lazy_tensor._solve(
                        torch.cat([probe_vectors, eager_rhs], -1), preconditioner=base_lazy_tensor._preconditioner()[0]
                    )
                    probe_vector_tmats = torch.tensor([])
                probe_vector_solves = all_solves[..., : probe_vectors.size(-1)].detach()
                solves = all_solves[..., probe_vectors.size(-1) :]

                return (
                    solves.detach(),
                    probe_vectors.detach(),
                    probe_vector_norms.detach(),
                    probe_vector_solves.detach(),
                    probe_vector_tmats.detach(),
                )

            else:
                # Compute solves
                if settings.fast_computations.log_prob.on():
                    solves = base_lazy_tensor._solve(eager_rhs, preconditioner=base_lazy_tensor._preconditioner()[0])
                else:
                    solves = base_lazy_tensor.cholesky()._cholesky_solve(eager_rhs)
                dtype = solves.dtype
                device = solves.device
                return (
                    solves.detach(),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                )
//  [python]
def get_dihedral(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from: 
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
        * c4: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return torch.atan2( ( (torch.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,  
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )
//  [python]
def get_dihedral(c1, c2, c3, c4):
    """ Returns the dihedral angle in radians.
        Will use atan2 formula from: 
        https://en.wikipedia.org/wiki/Dihedral_angle#In_polymer_physics
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
        * c4: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2
    u3 = c4 - c3

    return torch.atan2( ( (torch.linalg.norm(u2, dim=-1, keepdim=True) * u1) * torch.cross(u2,u3, dim=-1) ).sum(dim=-1) ,  
                        (  torch.cross(u1,u2, dim=-1) * torch.cross(u2, u3, dim=-1) ).sum(dim=-1) )
//  [python]
def get_angle(c1, c2, c3):
    """ Returns the angle in radians.
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2

    # dont use acos since norms involved. 
    # better use atan2 formula: atan2(cross, dot) from here: 
    # https://johnblackburne.blogspot.com/2012/05/angle-between-two-3d-vectors.html

    # add a minus since we want the angle in reversed order - sidechainnet issues
    return torch.atan2( torch.norm(torch.cross(u1,u2, dim=-1), dim=-1), 
                        -(u1*u2).sum(dim=-1) )
//  [python]
def get_angle(c1, c2, c3):
    """ Returns the angle in radians.
        Inputs: 
        * c1: (batch, 3) or (3,)
        * c2: (batch, 3) or (3,)
        * c3: (batch, 3) or (3,)
    """
    u1 = c2 - c1
    u2 = c3 - c2

    # dont use acos since norms involved. 
    # better use atan2 formula: atan2(cross, dot) from here: 
    # https://johnblackburne.blogspot.com/2012/05/angle-between-two-3d-vectors.html

    # add a minus since we want the angle in reversed order - sidechainnet issues
    return torch.atan2( torch.linalg.norm(torch.cross(u1,u2, dim=-1), dim=-1), 
                        -(u1*u2).sum(dim=-1) )
//  [python]
def forward(self, x_mag: torch.Tensor, y_mag: torch.Tensor):
        """Calculate forward propagation.

        Args:
            x_mag: Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag: Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            torch.Tensor: Spectral convergence loss value.
        """
        return torch.norm(y_mag - x_mag, p="fro") / (torch.norm(y_mag, p="fro") + self.epsilon)
//  [python]
def forward(self, x_mag: torch.Tensor, y_mag: torch.Tensor):
        """Calculate forward propagation.

        Args:
            x_mag: Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag: Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
        Returns:
            torch.Tensor: Spectral convergence loss value.
        """
        return torch.linalg.norm(y_mag - x_mag, ord="fro") / (torch.linalg.norm(y_mag, ord="fro") + self.epsilon)
//  [python]
def _get_pnorm_after_step(self, module):
        optim = torch.optim.SGD(module.parameters(), lr=0.01)
        loss = self._get_output(module).sum()
        loss.backward()
        optim.step()
        return torch.norm(torch.stack([p.detach().norm() for p in module.parameters()]))
//  [python]
def _get_pnorm_after_step(self, module):
        optim = torch.optim.SGD(module.parameters(), lr=0.01)
        loss = self._get_output(module).sum()
        loss.backward()
        optim.step()
        return torch.linalg.norm(torch.stack([p.detach().norm() for p in module.parameters()]))
//  [python]
def Norm(a):
    """
    Returns the matrix norm of the UniTensor. The input tensor should be untagged. 

    If the given UniTensor is a matrix (rank-2), matrix norm will be calculated. If the given UniTensor is a vector (rank-1), vector norm will be calculated. If the given UniTensor has more than 2 ranks, the vector norm will be appllied to last dimension. 


    Args:
        a :
            a UniTensor[untagged]

    Return:
        UniTensor, 0-rank (constant)

    """

    if isinstance(a,UniTensor):
        if a.is_symm :
            raise Exception("Norm","[ERROR] cannot operate Norm on a symmetry tensor")

        if a.braket is not None:
            raise Exception("Norm","[ERROR] Norm can only operate on untagged UniTensor")

        if a.rowrank != 1 or len(a.labels)!=2:
            raise Exception("Norm","[ERROR] the input UniTensor should be rank-2, untagged with rowrank=1")


        #tmp = torch.norm(a.Storage)
        #if len(tmp.shape) != 0:
        #    return UniTensor(bonds=[tor10.Bond(tmp.shape[i]) for i in range(len(tmp.shape))],rowrank=1,torch_tensor=tmp,check=False)
        #else:
        tmp = UniTensor(bonds=[],labels=[],rowrank=0,check=False)
        tmp._mac(torch_tensor = torch.norm(a.Storage))
        return tmp
    else:
        raise Exception("Norm(UniTensor)","[ERROR] Norm can only accept UniTensor")
//  [python]
def Norm(a):
    """
    Returns the matrix norm of the UniTensor. The input tensor should be untagged. 

    If the given UniTensor is a matrix (rank-2), matrix norm will be calculated. If the given UniTensor is a vector (rank-1), vector norm will be calculated. If the given UniTensor has more than 2 ranks, the vector norm will be appllied to last dimension. 


    Args:
        a :
            a UniTensor[untagged]

    Return:
        UniTensor, 0-rank (constant)

    """

    if isinstance(a,UniTensor):
        if a.is_symm :
            raise Exception("Norm","[ERROR] cannot operate Norm on a symmetry tensor")

        if a.braket is not None:
            raise Exception("Norm","[ERROR] Norm can only operate on untagged UniTensor")

        if a.rowrank != 1 or len(a.labels)!=2:
            raise Exception("Norm","[ERROR] the input UniTensor should be rank-2, untagged with rowrank=1")


        #tmp = torch.norm(a.Storage)
        #if len(tmp.shape) != 0:
        #    return UniTensor(bonds=[tor10.Bond(tmp.shape[i]) for i in range(len(tmp.shape))],rowrank=1,torch_tensor=tmp,check=False)
        #else:
        tmp = UniTensor(bonds=[],labels=[],rowrank=0,check=False)
        tmp._mac(torch_tensor = torch.linalg.norm(a.Storage))
        return tmp
    else:
        raise Exception("Norm(UniTensor)","[ERROR] Norm can only accept UniTensor")
//  [python]
def _l2_normalize(d):
    d_reshaped = d.view(d.shape[0], -1, *(1 for _ in range(d.dim() - 2)))
    d /= torch.norm(d_reshaped, 2, dim=1, keepdim=True) + 1e-8
    return d
//  [python]
def _l2_normalize(d):
    d_reshaped = d.view(d.shape[0], -1, *(1 for _ in range(d.dim() - 2)))
    d /= torch.linalg.norm(d_reshaped, 2, dim=1, keepdim=True) + 1e-8
    return d
//  [python]
def calculate_regularizer_loss(self):
        regularizer_loss = 0.0
        for name, param in self.to_regularize_parameters:
            regularizer_loss = regularizer_loss + torch.norm(param, p=2)
        return regularizer_loss
//  [python]
def calculate_regularizer_loss(self):
        regularizer_loss = 0.0
        for name, param in self.to_regularize_parameters:
            regularizer_loss = regularizer_loss + torch.linalg.norm(param, ord=2)
        return regularizer_loss
//  [python]
def interpolate_zs(zs, weight=0.7):
  """blend a grid of zs.

  interpolate a grid of random zs so that they are closer to each other
  helps to generate more coherent content

  Args:
    zs: torch.Tensor HxWx1xC
    weight: float

  Returns:
    blended latents
  """

  grid_size = zs.shape[0]
  zmid = zs[[grid_size // 2], [grid_size // 2], :, :]
  zinterp = zmid * weight + zs * (1 - weight)
  zinterp = (
      zinterp
      / torch.norm(zinterp, dim=-1, keepdim=True)
      * torch.norm(zs, dim=-1, keepdim=True)
  )
  zs = zinterp
  return zs
//  [python]
def interpolate_zs(zs, weight=0.7):
  """blend a grid of zs.

  interpolate a grid of random zs so that they are closer to each other
  helps to generate more coherent content

  Args:
    zs: torch.Tensor HxWx1xC
    weight: float

  Returns:
    blended latents
  """

  grid_size = zs.shape[0]
  zmid = zs[[grid_size // 2], [grid_size // 2], :, :]
  zinterp = zmid * weight + zs * (1 - weight)
  zinterp = (
      zinterp
      / torch.linalg.norm(zinterp, dim=-1, keepdim=True)
      * torch.linalg.norm(zs, dim=-1, keepdim=True)
  )
  zs = zinterp
  return zs
//  [python]
def getInvariant(self, positions):
        new_pos = torch.zeros(self.n_replicas, 1, self.n_particles*3)
        for i in range(len(positions)):
            inp = positions[i].view(-1, self.n_particles, 3)
            com = torch.mean(inp, dim=1)
            centered = inp-com[:,None,:]
            
            top = centered[:,-1]
            A = top / torch.norm(top, dim=1)[:, None]

            B = torch.tensor([1., 0., 0.], device=A.device)

            B = B.type(A.dtype)

            dot = (A @ B)
            cross = torch.cross(A, B.repeat(A.size(0), 1))
            cross_norm = torch.norm(cross, dim=1)

            zeros = torch.zeros_like(dot, device=dot.device)
            ones = torch.ones_like(dot, device=dot.device)
            G1 = torch.vstack((dot, -cross_norm, zeros)).T
            G2 = torch.vstack((cross_norm, dot, zeros)).T
            G3 = torch.vstack((zeros, zeros, ones)).T
            G = torch.stack((G1, G2, G3), dim=2)

            u = A
            v = (B - (dot[:, None] * A)) / torch.norm((B - (dot[:, None] * A)), dim=1)[:, None]
            w = torch.cross(B.repeat(A.size(0), 1), A)
            F_inv = torch.stack([u, v, w], dim=2)
            F = torch.inverse(F_inv)
            U = torch.matmul(torch.matmul(F_inv, G), F)
            rotated = torch.matmul(centered, U)
            new_pos[i] = rotated.view((1, self.n_particles*3))
        return new_pos
//  [python]
def getInvariant(self, positions):
        new_pos = torch.zeros(self.n_replicas, 1, self.n_particles*3)
        for i in range(len(positions)):
            inp = positions[i].view(-1, self.n_particles, 3)
            com = torch.mean(inp, dim=1)
            centered = inp-com[:,None,:]
            
            top = centered[:,-1]
            A = top / torch.linalg.norm(top, dim=1)[:, None]

            B = torch.tensor([1., 0., 0.], device=A.device)

            B = B.type(A.dtype)

            dot = (A @ B)
            cross = torch.cross(A, B.repeat(A.size(0), 1))
            cross_norm = torch.linalg.norm(cross, dim=1)

            zeros = torch.zeros_like(dot, device=dot.device)
            ones = torch.ones_like(dot, device=dot.device)
            G1 = torch.vstack((dot, -cross_norm, zeros)).T
            G2 = torch.vstack((cross_norm, dot, zeros)).T
            G3 = torch.vstack((zeros, zeros, ones)).T
            G = torch.stack((G1, G2, G3), dim=2)

            u = A
            v = (B - (dot[:, None] * A)) / torch.linalg.norm((B - (dot[:, None] * A)), dim=1)[:, None]
            w = torch.cross(B.repeat(A.size(0), 1), A)
            F_inv = torch.stack([u, v, w], dim=2)
            F = torch.inverse(F_inv)
            U = torch.matmul(torch.matmul(F_inv, G), F)
            rotated = torch.matmul(centered, U)
            new_pos[i] = rotated.view((1, self.n_particles*3))
        return new_pos
//  [python]
def get_grad_norm_(parameters, norm_type: float = 2.0, layer_names=None) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    
    parameters = [p for p in parameters if p.grad is not None]
        
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        # total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
        layer_norm = torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters])
        total_norm = torch.norm(layer_norm, norm_type)
        # print(layer_norm.max(dim=0))
        
        if layer_names is not None:
            if torch.isnan(total_norm) or torch.isinf(total_norm) or total_norm > 1.0:
                value_top, name_top = torch.topk(layer_norm, k=5)
                print(f"Top norm value: {value_top}")
                print(f"Top norm name: {[layer_names[i][7:] for i in name_top.tolist()]}")
        
    return total_norm
//  [python]
def get_grad_norm_(parameters, norm_type: float = 2.0, layer_names=None) -> torch.Tensor:
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    
    parameters = [p for p in parameters if p.grad is not None]
        
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    
    if norm_type == inf:
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        # total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
        layer_norm = torch.stack([torch.linalg.norm(p.grad.detach(), norm_type).to(device) for p in parameters])
        total_norm = torch.linalg.norm(layer_norm, norm_type)
        # print(layer_norm.max(dim=0))
        
        if layer_names is not None:
            if torch.isnan(total_norm) or torch.isinf(total_norm) or total_norm > 1.0:
                value_top, name_top = torch.topk(layer_norm, k=5)
                print(f"Top norm value: {value_top}")
                print(f"Top norm name: {[layer_names[i][7:] for i in name_top.tolist()]}")
        
    return total_norm
//  [python]
def pairwise_cosine_matrix(matrix1, matrix2):
    dot = torch.matmul(matrix1, matrix2.transpose(1, 2))
    matrix1_norm = torch.norm(matrix1, dim=-1, keepdim=True)
    matrix2_norm = torch.norm(matrix2, dim=-1, keepdim=True)
    norm = torch.matmul(matrix1_norm, matrix2_norm.transpose(1, 2))
    return dot / norm
//  [python]
def pairwise_cosine_matrix(matrix1, matrix2):
    dot = torch.matmul(matrix1, matrix2.transpose(1, 2))
    matrix1_norm = torch.linalg.norm(matrix1, dim=-1, keepdim=True)
    matrix2_norm = torch.linalg.norm(matrix2, dim=-1, keepdim=True)
    norm = torch.matmul(matrix1_norm, matrix2_norm.transpose(1, 2))
    return dot / norm
//  [python]
def forward(self, x, y, dij=None):
        x = LogEigFunction.apply(x)
        x = x.view(x.size(0), -1)
        y = LogEigFunction.apply(y)
        y = y.view(y.size(0), -1)
        fe = pair(y, x, dij)
        if check_gpu():
            fe = fe.cuda()

        x1 = fe[:, : x.size(1)].contiguous()
        x2 = fe[:, x.size(1) :].contiguous()
        d = torch.norm(x1 - x2 + 1e-16, dim=1)
        return d
//  [python]
def forward(self, x, y, dij=None):
        x = LogEigFunction.apply(x)
        x = x.view(x.size(0), -1)
        y = LogEigFunction.apply(y)
        y = y.view(y.size(0), -1)
        fe = pair(y, x, dij)
        if check_gpu():
            fe = fe.cuda()

        x1 = fe[:, : x.size(1)].contiguous()
        x2 = fe[:, x.size(1) :].contiguous()
        d = torch.linalg.norm(x1 - x2 + 1e-16, dim=1)
        return d
//  [python]
def forward(self, x, y, dij=None):
        x = LogEigFunction.apply(x)
        x = x.view(x.size(0), -1)
        y = LogEigFunction.apply(y)
        y = y.view(y.size(0), -1)
        fe = pair(y, x, dij)
        if check_gpu():
            fe = fe.cuda()

        x1 = fe[:, : x.size(1)].contiguous()
        x1 = self.dimLinear(x1)
        x2 = fe[:, x.size(1) :].contiguous()
        x2 = self.dimLinear(x2)
        d = torch.norm(x1 - x2 + 1e-16, dim=1)
        return d
//  [python]
def forward(self, x, y, dij=None):
        x = LogEigFunction.apply(x)
        x = x.view(x.size(0), -1)
        y = LogEigFunction.apply(y)
        y = y.view(y.size(0), -1)
        fe = pair(y, x, dij)
        if check_gpu():
            fe = fe.cuda()

        x1 = fe[:, : x.size(1)].contiguous()
        x1 = self.dimLinear(x1)
        x2 = fe[:, x.size(1) :].contiguous()
        x2 = self.dimLinear(x2)
        d = torch.linalg.norm(x1 - x2 + 1e-16, dim=1)
        return d
//  [python]
def axis_angle_add(theta, roll_axis, alpha):
    """Composition of two axis-angle rotations (PyTorch version)
    Args:
        theta: N x 3
        roll_axis: N x 3
        alph: N x 1
    Returns:
        equivalent axis-angle of the composition
    """
    l2norm = torch.norm(theta + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(l2norm, -1)

    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    b_cos = torch.cos(angle).cpu()
    b_sin = torch.sin(angle).cpu()

    a_cos = torch.cos(alpha)
    a_sin = torch.sin(alpha)

    dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True)
    cross_mm = torch.zeros_like(normalized)
    cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
    cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
    cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]

    c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
    c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm

    c_angle = 2 * torch.acos(c_cos)
    c_sin = torch.sin(c_angle * 0.5)
    c_n = (c_angle / c_sin) * c_sin_n

    return c_n
//  [python]
def axis_angle_add(theta, roll_axis, alpha):
    """Composition of two axis-angle rotations (PyTorch version)
    Args:
        theta: N x 3
        roll_axis: N x 3
        alph: N x 1
    Returns:
        equivalent axis-angle of the composition
    """
    l2norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l2norm, -1)

    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    b_cos = torch.cos(angle).cpu()
    b_sin = torch.sin(angle).cpu()

    a_cos = torch.cos(alpha)
    a_sin = torch.sin(alpha)

    dot_mm = torch.sum(normalized * roll_axis, dim=1, keepdim=True)
    cross_mm = torch.zeros_like(normalized)
    cross_mm[:, 0] = roll_axis[:, 1] * normalized[:, 2] - roll_axis[:, 2] * normalized[:, 1]
    cross_mm[:, 1] = roll_axis[:, 2] * normalized[:, 0] - roll_axis[:, 0] * normalized[:, 2]
    cross_mm[:, 2] = roll_axis[:, 0] * normalized[:, 1] - roll_axis[:, 1] * normalized[:, 0]

    c_cos = a_cos * b_cos - a_sin * b_sin * dot_mm
    c_sin_n = a_sin * b_cos * roll_axis + a_cos * b_sin * normalized + a_sin * b_sin * cross_mm

    c_angle = 2 * torch.acos(c_cos)
    c_sin = torch.sin(c_angle * 0.5)
    c_n = (c_angle / c_sin) * c_sin_n

    return c_n
//  [python]
def _fit_projection(self, X, d=None):
        if d is None:
            d = self.d
        if d > 1:
            raise ValueError
        N = X.shape[1]
        V_init = init_coef(N, d, self.rng, self.init)

        v = torch.tensor(V_init, requires_grad=True,
                         device=self.device, dtype=self.dtype)

        Xt = X
        if not isinstance(Xt, torch.Tensor):
            Xt = torch.tensor(Xt, device=self.device, dtype=self.dtype)

        if self.verbose:
            def callback(v_flat):
                with torch.no_grad():
                    v_flat_torch = torch.tensor(v_flat,
                                                device=self.device,
                                                dtype=self.dtype)
                    v_torch = v_flat_torch.reshape(N, d)
                    v_torch = v_torch / torch.norm(v_torch, dim=0, keepdim=True)
                    ent = ent_loss_fn(Xt, v_torch, self.T)
                    ent = ent.detach().cpu().numpy()
                    print("Ent: {} bits".format(str(np.round(ent, 4))))
            callback(V_init)
        else:
            callback = None

        def f_df(v_flat):
            v_flat_torch = torch.tensor(v_flat,
                                        requires_grad=True,
                                        device=self.device,
                                        dtype=self.dtype)
            v_torch = v_flat_torch.reshape(N, d)
            v_torch = v_torch / torch.norm(v_torch, dim=0, keepdim=True)
            loss = ent_loss_fn(Xt, v_torch, self.T)
            loss.backward()
            grad = v_flat_torch.grad
            return (loss.detach().cpu().numpy().astype(float),
                    grad.detach().cpu().numpy().astype(float))
        opt = minimize(f_df, V_init.ravel(), method='L-BFGS-B', jac=True,
                       options={'disp': self.verbose, 'ftol': 1e-6, 'gtol': 1e-5,
                                'maxfun': 15000, 'maxiter': 15000, 'maxls': 20},
                       callback=callback)
        v = opt.x.reshape(N, d)

        # Orthonormalize the basis prior to returning it
        V_opt = scipy.linalg.orth(v)
        with torch.no_grad():
            v_flat_torch = torch.tensor(V_opt.ravel(),
                                        device=self.device,
                                        dtype=self.dtype)
            v_torch = v_flat_torch.reshape(N, d)
            final_pi = ent_loss_fn(Xt, v_torch, self.T).detach().cpu().numpy()
        return V_opt, final_pi
//  [python]
def _fit_projection(self, X, d=None):
        if d is None:
            d = self.d
        if d > 1:
            raise ValueError
        N = X.shape[1]
        V_init = init_coef(N, d, self.rng, self.init)

        v = torch.tensor(V_init, requires_grad=True,
                         device=self.device, dtype=self.dtype)

        Xt = X
        if not isinstance(Xt, torch.Tensor):
            Xt = torch.tensor(Xt, device=self.device, dtype=self.dtype)

        if self.verbose:
            def callback(v_flat):
                with torch.no_grad():
                    v_flat_torch = torch.tensor(v_flat,
                                                device=self.device,
                                                dtype=self.dtype)
                    v_torch = v_flat_torch.reshape(N, d)
                    v_torch = v_torch / torch.linalg.norm(v_torch, dim=0, keepdim=True)
                    ent = ent_loss_fn(Xt, v_torch, self.T)
                    ent = ent.detach().cpu().numpy()
                    print("Ent: {} bits".format(str(np.round(ent, 4))))
            callback(V_init)
        else:
            callback = None

        def f_df(v_flat):
            v_flat_torch = torch.tensor(v_flat,
                                        requires_grad=True,
                                        device=self.device,
                                        dtype=self.dtype)
            v_torch = v_flat_torch.reshape(N, d)
            v_torch = v_torch / torch.linalg.norm(v_torch, dim=0, keepdim=True)
            loss = ent_loss_fn(Xt, v_torch, self.T)
            loss.backward()
            grad = v_flat_torch.grad
            return (loss.detach().cpu().numpy().astype(float),
                    grad.detach().cpu().numpy().astype(float))
        opt = minimize(f_df, V_init.ravel(), method='L-BFGS-B', jac=True,
                       options={'disp': self.verbose, 'ftol': 1e-6, 'gtol': 1e-5,
                                'maxfun': 15000, 'maxiter': 15000, 'maxls': 20},
                       callback=callback)
        v = opt.x.reshape(N, d)

        # Orthonormalize the basis prior to returning it
        V_opt = scipy.linalg.orth(v)
        with torch.no_grad():
            v_flat_torch = torch.tensor(V_opt.ravel(),
                                        device=self.device,
                                        dtype=self.dtype)
            v_torch = v_flat_torch.reshape(N, d)
            final_pi = ent_loss_fn(Xt, v_torch, self.T).detach().cpu().numpy()
        return V_opt, final_pi
//  [python]
def sphere_plot(
    ax, 
    samples
):
    x = samples
    x = x / torch.norm(x, dim=-1, keepdim=True)
    angle = torch.zeros(*x.shape[:-1], 2)
    angle[..., 0] = torch.atan2(x[..., 1], x[..., 0])
    angle[..., 1] = torch.arcsin(x[..., 2])
    samples = (180 * angle / np.pi).reshape(-1,2)

    lon, lat = samples[:, 0].cpu().numpy(), samples[:, 1].cpu().numpy()
    basemap = Basemap(projection='robin', lat_0=0, lon_0=0)
    x, y = basemap(lon, lat)
        
    continents_color = 'lightcyan'
    water_color = 'deepskyblue'
    size=10
    alpha=0.1
    marker="o"
    color='crimson'
    
    basemap.fillcontinents(color=continents_color, lake_color=water_color)
    basemap.drawmapboundary(fill_color=water_color)
    ax.scatter(
        x, y, 
        s=size,
        alpha=alpha,
        c=color, 
        marker=marker
    )
    pass
//  [python]
def sphere_plot(
    ax, 
    samples
):
    x = samples
    x = x / torch.linalg.norm(x, dim=-1, keepdim=True)
    angle = torch.zeros(*x.shape[:-1], 2)
    angle[..., 0] = torch.atan2(x[..., 1], x[..., 0])
    angle[..., 1] = torch.arcsin(x[..., 2])
    samples = (180 * angle / np.pi).reshape(-1,2)

    lon, lat = samples[:, 0].cpu().numpy(), samples[:, 1].cpu().numpy()
    basemap = Basemap(projection='robin', lat_0=0, lon_0=0)
    x, y = basemap(lon, lat)
        
    continents_color = 'lightcyan'
    water_color = 'deepskyblue'
    size=10
    alpha=0.1
    marker="o"
    color='crimson'
    
    basemap.fillcontinents(color=continents_color, lake_color=water_color)
    basemap.drawmapboundary(fill_color=water_color)
    ax.scatter(
        x, y, 
        s=size,
        alpha=alpha,
        c=color, 
        marker=marker
    )
    pass
//  [python]
def wav_to_energy(y, n_fft, hop_length, win_length, center=False):
    spec = _wav_to_spec(y, n_fft, hop_length, win_length, center=center)

    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
    return torch.norm(spec, dim=1, keepdim=True)
//  [python]
def wav_to_energy(y, n_fft, hop_length, win_length, center=False):
    spec = _wav_to_spec(y, n_fft, hop_length, win_length, center=center)

    spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
    return torch.linalg.norm(spec, dim=1, keepdim=True)
//  [python]
def forward(self, h_head, h_tail, rels):
        r"""

        Description
        -----------
        Score triples.

        Parameters
        ----------
        h_head : torch.Tensor
            Head entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        h_tail : torch.Tensor
            Tail entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        rels : torch.Tensor
            Relation types. It is a LongTensor of shape :math:`(E)`, where
            :math:`E` is the number of triples.

        Returns
        -------
        torch.Tensor
            The triple scores. The tensor is of shape :math:`(E)`.
        """
        h_rel = self.rel_emb(rels)

        return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)
//  [python]
def forward(self, h_head, h_tail, rels):
        r"""

        Description
        -----------
        Score triples.

        Parameters
        ----------
        h_head : torch.Tensor
            Head entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        h_tail : torch.Tensor
            Tail entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        rels : torch.Tensor
            Relation types. It is a LongTensor of shape :math:`(E)`, where
            :math:`E` is the number of triples.

        Returns
        -------
        torch.Tensor
            The triple scores. The tensor is of shape :math:`(E)`.
        """
        h_rel = self.rel_emb(rels)

        return -torch.linalg.norm(h_head + h_rel - h_tail, ord=self.p, dim=-1)
//  [python]
def g(F):
        FF = F.t() @ F
        FF_norm = torch.norm(FF, p='fro')
        return (1 / (FF_norm + epsilon_F)) * FF
//  [python]
def g(F):
        FF = F.t() @ F
        FF_norm = torch.linalg.norm(FF, ord='fro')
        return (1 / (FF_norm + epsilon_F)) * FF
//  [python]
def test_iconv(self, channels=64, seed=2019,h=8):
        torch.random.manual_seed(seed)

        weight_obj = iConv2d(channels, channels)
        w=h
        input_activation = torch.randn(1,channels,h,w)
        _ = weight_obj(input_activation)
        weight = weight_obj.conv.weight
        weight_numpy = weight.detach().cpu().permute((2,3,0,1)).numpy()

        # compute 2d fft 
       # print(weight_numpy.shape)
        kernel_fft = np.fft.fft2(weight_numpy,[h,w], axes=[0,1], norm=None)
        padded_numpy = np.pad(weight_numpy,((0,h-3),(0,w-3),(0,0),(0,0)),mode='constant')
        kernel_fft2 = np.fft.fft2(padded_numpy, axes=[0,1])
        #print("original",(kernel_fft-kernel_fft2))
        # then take svds
        svds = np.linalg.svd(kernel_fft, compute_uv=False)
        # finally log det is sum(log(singular values))
        true_logdet = np.sum(np.log(svds))
        #print(np.min(svds))
        relative_error = torch.norm(true_logdet - weight_obj.logdet()) / np.linalg.norm(true_logdet)
        #print('relative error is: ', relative_error)
        self.assertLess(relative_error, 4e-3)
//  [python]
def test_iconv(self, channels=64, seed=2019,h=8):
        torch.random.manual_seed(seed)

        weight_obj = iConv2d(channels, channels)
        w=h
        input_activation = torch.randn(1,channels,h,w)
        _ = weight_obj(input_activation)
        weight = weight_obj.conv.weight
        weight_numpy = weight.detach().cpu().permute((2,3,0,1)).numpy()

        # compute 2d fft 
       # print(weight_numpy.shape)
        kernel_fft = np.fft.fft2(weight_numpy,[h,w], axes=[0,1], norm=None)
        padded_numpy = np.pad(weight_numpy,((0,h-3),(0,w-3),(0,0),(0,0)),mode='constant')
        kernel_fft2 = np.fft.fft2(padded_numpy, axes=[0,1])
        #print("original",(kernel_fft-kernel_fft2))
        # then take svds
        svds = np.linalg.svd(kernel_fft, compute_uv=False)
        # finally log det is sum(log(singular values))
        true_logdet = np.sum(np.log(svds))
        #print(np.min(svds))
        relative_error = torch.linalg.norm(true_logdet - weight_obj.logdet()) / np.linalg.norm(true_logdet)
        #print('relative error is: ', relative_error)
        self.assertLess(relative_error, 4e-3)
//  [python]
def _safe_normalize(x, threshold=None):
    norm = torch.norm(x)
    if threshold is None:
        threshold = torch.finfo(norm.dtype).eps
    normalized_x = x / norm if norm > threshold else torch.zeros_like(x)
    return normalized_x, norm
//  [python]
def _safe_normalize(x, threshold=None):
    norm = torch.linalg.norm(x)
    if threshold is None:
        threshold = torch.finfo(norm.dtype).eps
    normalized_x = x / norm if norm > threshold else torch.zeros_like(x)
    return normalized_x, norm
//  [python]
def train_nn_lyapunov_approximator(lyapunov_relu, R, target_lyapunov_relu,
                                   target_R, V_lambda, x_equilibrium, x_lo,
                                   x_up, num_samples, num_epochs):
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    with torch.no_grad():
        target_V = target_lyapunov_relu(x_samples) - target_lyapunov_relu(
            x_equilibrium) + V_lambda * torch.norm(
                target_R @ (x_samples - x_equilibrium).T, p=1, dim=0).reshape(
                    (-1, 1))

    dataset = torch.utils.data.TensorDataset(x_samples, target_V)

    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))

    R.requires_grad_(True)
    utils.train_approximator(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_nn_lyapunov_approximator(lyapunov_relu, R, target_lyapunov_relu,
                                   target_R, V_lambda, x_equilibrium, x_lo,
                                   x_up, num_samples, num_epochs):
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    with torch.no_grad():
        target_V = target_lyapunov_relu(x_samples) - target_lyapunov_relu(
            x_equilibrium) + V_lambda * torch.linalg.norm(
                target_R @ (x_samples - x_equilibrium).T, ord=1, dim=0).reshape(
                    (-1, 1))

    dataset = torch.utils.data.TensorDataset(x_samples, target_V)

    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))

    R.requires_grad_(True)
    utils.train_approximator(dataset,
                             lyapunov_relu,
                             compute_V,
                             batch_size=50,
                             num_epochs=num_epochs,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def k_svd(X, M, s, n_iter, max_time=None):
    assert n_iter >= 1
    n, d = X.shape
    k, _ = M.shape

    start = time.time()
    last_error = 10**10
    for iter in range(n_iter):
        # Sparse Coding
        ids, m = omp(X, M, s)

        # Dictionary Update, one row at a time
        for j in range(k):
            # Find which samples are currently assumed to be using the kth atom
            mask = ids == j
            I = torch.any(mask, dim=1)

            if not torch.any(I):
                continue
            # Compute current residuals
            E = X[I] - (m[I].unsqueeze(1) @ M[ids[I]]).squeeze(1)  # E = X[I] - S[I] @ M
            # Add the current atom back in. This is like assuming it was
            # currently not used by any of the samples.
            E += torch.outer(m[mask], M[j])  # E += torch.outer(S[I, j], M[j])

            # Use svd_lowrank from torch to save time and only compute
            # the first singular vector
            U, Sigma, V = torch.svd_lowrank(E, q=1)
            M[j, :] = V[:, 0]
            # We also update S, which is how k-svd can have an advantage over MOD
            # S[I, j] = Sigma[0] * U[:, 0]
            m[mask] = Sigma[0] * U[:, 0]

        SM = (m.unsqueeze(1) @ M[ids]).squeeze(1)  # SM = S @ M
        error = torch.norm(SM - X)
        if error < 1e-4:
            print("K-SVD: Stopping early because error is near 0.")
            break

        if error > last_error:
            print(f"K-SVD: Stopping early because error is growing. last_error={last_error.item()}.")
            break
        last_error = error

        if max_time is not None and time.time() - start > max_time:
            print("K-SVD: Stopping early because ran out of time.")
            break

    print(f"K-SVD: error at {iter=}:", error.item())

    return M, ids, m
//  [python]
def k_svd(X, M, s, n_iter, max_time=None):
    assert n_iter >= 1
    n, d = X.shape
    k, _ = M.shape

    start = time.time()
    last_error = 10**10
    for iter in range(n_iter):
        # Sparse Coding
        ids, m = omp(X, M, s)

        # Dictionary Update, one row at a time
        for j in range(k):
            # Find which samples are currently assumed to be using the kth atom
            mask = ids == j
            I = torch.any(mask, dim=1)

            if not torch.any(I):
                continue
            # Compute current residuals
            E = X[I] - (m[I].unsqueeze(1) @ M[ids[I]]).squeeze(1)  # E = X[I] - S[I] @ M
            # Add the current atom back in. This is like assuming it was
            # currently not used by any of the samples.
            E += torch.outer(m[mask], M[j])  # E += torch.outer(S[I, j], M[j])

            # Use svd_lowrank from torch to save time and only compute
            # the first singular vector
            U, Sigma, V = torch.svd_lowrank(E, q=1)
            M[j, :] = V[:, 0]
            # We also update S, which is how k-svd can have an advantage over MOD
            # S[I, j] = Sigma[0] * U[:, 0]
            m[mask] = Sigma[0] * U[:, 0]

        SM = (m.unsqueeze(1) @ M[ids]).squeeze(1)  # SM = S @ M
        error = torch.linalg.norm(SM - X)
        if error < 1e-4:
            print("K-SVD: Stopping early because error is near 0.")
            break

        if error > last_error:
            print(f"K-SVD: Stopping early because error is growing. last_error={last_error.item()}.")
            break
        last_error = error

        if max_time is not None and time.time() - start > max_time:
            print("K-SVD: Stopping early because ran out of time.")
            break

    print(f"K-SVD: error at {iter=}:", error.item())

    return M, ids, m
//  [python]
def mini_ksvd(M, s, batch_maker, n_batches, n_iter, max_time=None):
    # Inspired by https://csaws.cs.technion.ac.il/~ronrubin/Publications/KSVD-OMP-v2.pdf
    assert n_iter >= 1
    k, _ = M.shape

    start = time.time()
    last_error = 10**10
    for iter in range(n_iter):
        error = 0
        new_Ms = torch.zeros_like(M)
        n = 0
        for j in tqdm(range(n_batches)):
            X = batch_maker(j)
            if X is None:
                break
            new_M = M.clone()
            # Sparse coding
            ids, m = omp(X, new_M, s)
            # Find new dictionary
            for j in range(k):
                mask = ids == j
                I = torch.any(mask, dim=1)
                if not torch.any(I):
                    continue
                E = X[I] - (m[I].unsqueeze(1) @ new_M[ids[I]]).squeeze(1) + torch.outer(m[mask], new_M[j])

                # Could use the cheaper alternative decomposition in Approximate K-SVD here
                U, Sigma, V = torch.svd_lowrank(E, q=1)
                new_M[j, :] = V[:, 0]
                m[mask] = Sigma[0] * U[:, 0]

            # "Gradient" update
            # M = (1-lr)*M + lr*new_M
            # M = new_M
            new_Ms += new_M
            n += 1

            SM = (m.unsqueeze(1) @ M[ids]).squeeze(1)  # SM = S @ M
            error += torch.norm(SM - X) ** 2

        M = new_Ms / n

        error = error**0.5
        if error < 1e-4:
            print("K-SVD: Stopping early because error is near 0.")
            break

        if error > last_error:
            print(f"K-SVD: Stopping early because error is growing. last_error={last_error.item()}.")
            break
        last_error = error

        if max_time is not None and time.time() - start > max_time:
            print("K-SVD: Stopping early because ran out of time.")
            break

    print(f"K-SVD: error at {iter=}:", error.item())

    return M