jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def becke_tesellation(
            self, x: torch.Tensor, r: torch.Tensor, dm: torch.Tensor,
            i: torch.Tensor
    ) -> torch.Tensor:
        """

        Generates a weight scheme to avoid overlap

        :param x: sample
        :param r: centers
        :param dm: centers distance matrix
        :param i: map sample -> center
        :return:
        """
        device = x.device
        dtype = x.dtype

        nx = x.size()[0]  # number of sample matches. It should match nx
        mx = x.size()[1]  # number of samples per molecule
        nr = r.size()[0]  # number of molecules
        mr = r.size()[1]  # number of atoms per molecule
        if nx != nr:
            raise IOError("batch dims don't not match: {:d} {:d}".format(nx, nr))

        # calculate r-r distance matrix

        # calculate x-r distance matrix
        # dim 1 will be r and dim 2 will be x
        # final dims should be [nr, mx, mr]

        x_exp = x.unsqueeze(2).expand(nr, mx, mr, 3)
        r_exp = r.unsqueeze(1).expand(nr, mx, mr, 3)
        rx_dm = torch.linalg.norm(x_exp - r_exp, dim=3)
        del x_exp, r_exp

        # calculate (x-r)-(x-r)T distance vectors

        rx_dm_expanded1 = rx_dm.unsqueeze(3).expand(nr, mx, mr, mr)
        rx_dm_expanded2 = rx_dm.unsqueeze(2).expand(nr, mx, mr, mr)
        xx_dm_expanded = dm.unsqueeze(0)
        mu = (rx_dm_expanded1 - rx_dm_expanded2) / xx_dm_expanded.clamp(min=1e-12)

        del rx_dm_expanded1, rx_dm_expanded2, xx_dm_expanded
        del dm, rx_dm
        # soft-cutoff function

        for _ in range(1, self.softening):
            mu = 0.5 * mu * (3.0 - (mu.pow(2.0)))

        mu = 0.5 * (1 - mu)

        # diagonal fill

        diag_fill = torch.eye(
            mr, mr, device=device, dtype=dtype
        ).unsqueeze(0).unsqueeze(1).expand(nr, mx, mr, mr)

        mu += (diag_fill * 0.5)

        # productory

        mu = mu.prod(dim=3)

        # mu should be now [nr, mx, mr] containing
        # the weights of each nuclei for each sampling point

        i = i.unsqueeze(2)

        v_i = torch.gather(mu, 2, i).squeeze(2)
        v_all = mu.sum(2)
        w = v_i / v_all

        # w should be now [nr, mx] containing the weight of each sample

        return w
//  [python]
def dist(self, x, y):
        """Batched geodesic distance"""
        diff = torch.bmm(x, self.inv(y))
        torus_diff = self.torus_representative(diff)
        log_torus_diff = torch.arccos(torus_diff.real)
        dist = math.sqrt(2) * torch.minimum(log_torus_diff, 2 * math.pi - log_torus_diff)
        dist = torch.norm(dist, dim=1)
        return dist
//  [python]
def dist(self, x, y):
        """Batched geodesic distance"""
        diff = torch.bmm(x, self.inv(y))
        torus_diff = self.torus_representative(diff)
        log_torus_diff = torch.arccos(torus_diff.real)
        dist = math.sqrt(2) * torch.minimum(log_torus_diff, 2 * math.pi - log_torus_diff)
        dist = torch.linalg.norm(dist, dim=1)
        return dist
//  [python]
def pairwise_dist(self, x, y):
        """For n points x_i and m points y_j computed dist(x_i,y_j)"""
        x_y_ = self.pairwise_embed(x, y)
        log_x_y_ = torch.arccos(x_y_.real)
        dist = math.sqrt(2) * torch.minimum(log_x_y_, 2 * math.pi - log_x_y_)
        dist = torch.norm(dist, dim=1).reshape(x.shape[0], y.shape[0])
        return dist
//  [python]
def pairwise_dist(self, x, y):
        """For n points x_i and m points y_j computed dist(x_i,y_j)"""
        x_y_ = self.pairwise_embed(x, y)
        log_x_y_ = torch.arccos(x_y_.real)
        dist = math.sqrt(2) * torch.minimum(log_x_y_, 2 * math.pi - log_x_y_)
        dist = torch.linalg.norm(dist, dim=1).reshape(x.shape[0], y.shape[0])
        return dist
//  [python]
def __exit__(self, exc_type, exc_value, exc_traceback):
        # Dump out all model gnorms to enable better debugging
        norm = {}
        gradients = {}
        for name, param in self.named_parameters:
            if param.grad is not None:
                grad_norm = torch.norm(param.grad.data, p=2, dtype=torch.float32)
                norm[name] = grad_norm.item()
                if torch.isnan(grad_norm).any() or torch.isinf(grad_norm).any():
                    gradients[name] = param.grad.data
        if len(gradients) > 0:
            logger.info("Detected nan/inf grad norm, dumping norms...")
            logger.info(f"norms: {norm}")
            logger.info(f"gradients: {gradients}")

        self.close()
//  [python]
def __exit__(self, exc_type, exc_value, exc_traceback):
        # Dump out all model gnorms to enable better debugging
        norm = {}
        gradients = {}
        for name, param in self.named_parameters:
            if param.grad is not None:
                grad_norm = torch.linalg.norm(param.grad.data, ord=2, dtype=torch.float32)
                norm[name] = grad_norm.item()
                if torch.isnan(grad_norm).any() or torch.isinf(grad_norm).any():
                    gradients[name] = param.grad.data
        if len(gradients) > 0:
            logger.info("Detected nan/inf grad norm, dumping norms...")
            logger.info(f"norms: {norm}")
            logger.info(f"gradients: {gradients}")

        self.close()
//  [python]
def print_test(cmd_opt1: torch.Tensor, cmd_opt2: torch.Tensor, grad: torch.Tensor,
               val1: torch.Tensor = None, val2: torch.Tensor = None):
    """
    Check if solutions are close
    """
    error = torch.norm(cmd_opt1 - cmd_opt2) / torch.norm(cmd_opt1)
    # print('Precision machine is at 1e-16, precision of the algorithms depend then on the norm of the gradient')
    print('Norm gradient by auto-diff : {0:.2e}'.format(torch.norm(grad)))
    print('Relative norm difference solutions: {0:.2e} All close ? {1}'.format(error, cmd_opt1.allclose(cmd_opt2)))
    print('Relative difference optimal value: {0:.2e}\n'.format(torch.abs(val1 - val2)/torch.abs(val1)))
//  [python]
def print_test(cmd_opt1: torch.Tensor, cmd_opt2: torch.Tensor, grad: torch.Tensor,
               val1: torch.Tensor = None, val2: torch.Tensor = None):
    """
    Check if solutions are close
    """
    error = torch.linalg.norm(cmd_opt1 - cmd_opt2) / torch.linalg.norm(cmd_opt1)
    # print('Precision machine is at 1e-16, precision of the algorithms depend then on the norm of the gradient')
    print('Norm gradient by auto-diff : {0:.2e}'.format(torch.linalg.norm(grad)))
    print('Relative norm difference solutions: {0:.2e} All close ? {1}'.format(error, cmd_opt1.allclose(cmd_opt2)))
    print('Relative difference optimal value: {0:.2e}\n'.format(torch.abs(val1 - val2)/torch.abs(val1)))
//  [python]
def compute_weights(raw, z_vals, rays_d, noise=0.):
    raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists)

    dists = z_vals[...,1:] - z_vals[...,:-1]
    dists = torch.cat([dists, torch.full_like(dists[...,:1], 1e10, device=device)], -1)  # [N_rays, N_samples]
    dists = dists * torch.norm(rays_d[...,None,:], dim=-1)
    
    alpha = raw2alpha(raw[...,3] + noise, dists)  # [N_rays, N_samples]
    # weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
    weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1), device=device), 1.-alpha + 1e-10], -1), -1)[:, :-1]

    return weights
//  [python]
def compute_weights(raw, z_vals, rays_d, noise=0.):
    raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists)

    dists = z_vals[...,1:] - z_vals[...,:-1]
    dists = torch.cat([dists, torch.full_like(dists[...,:1], 1e10, device=device)], -1)  # [N_rays, N_samples]
    dists = dists * torch.linalg.norm(rays_d[...,None,:], dim=-1)
    
    alpha = raw2alpha(raw[...,3] + noise, dists)  # [N_rays, N_samples]
    # weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
    weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1), device=device), 1.-alpha + 1e-10], -1), -1)[:, :-1]

    return weights
//  [python]
def step(self, closure=None):
    hps = self.hps
    for group in self.param_groups:
      lr = group['lr']
      for p in group['params']:
        if p.grad is None: continue
        grad = p.grad.data
        if grad.is_sparse:
          raise RuntimeError('Shampoo does not support sparse yet')
        state = self.state[p]
        if not state:
          self.init_var_state(p, state)
        state[STEP] += 1

        preconditioner = state[PRECONDITIONER]
        graft = state[GRAFT]

        # Gather statistics, compute preconditioners
        graft.add_statistics(grad)
        if state[STEP] % hps.statistics_compute_steps == 0:
          preconditioner.add_statistics(grad)
        if state[STEP] % hps.preconditioning_compute_steps == 0:
          preconditioner.compute_preconditioners()

        # Precondition gradients
        graft_grad = graft.precondition_gradient(grad)
        shampoo_grad = grad
        if state[STEP] >= self.hps.start_preconditioning_step:
          shampoo_grad = preconditioner.preconditioned_grad(grad)

        # Grafting
        graft_norm = torch.norm(graft_grad)
        shampoo_norm = torch.norm(shampoo_grad)
        shampoo_grad.mul_(graft_norm / (shampoo_norm + 1e-16))

        # Weight decay
        if self.hps.weight_decay != 0.0:
          shampoo_grad.add_(p.data, alpha=self.hps.weight_decay)
          graft_grad.add_(p.data, alpha=self.hps.weight_decay)

        # Momentum and Nesterov momentum, if needed
        state[MOMENTUM].mul_(group['momentum']).add_(shampoo_grad)
        graft_momentum = graft.update_momentum(grad, group['momentum'])

        if state[STEP] >= self.hps.start_preconditioning_step:
          momentum_update = state[MOMENTUM]
          wd_update = shampoo_grad
        else:
          momentum_update = graft_momentum
          wd_update = graft_grad

        if hps.nesterov:
          momentum_update.mul_(group['momentum']).add_(wd_update)

        # Final update
        p.data.add_(momentum_update, alpha=-lr)
//  [python]
def step(self, closure=None):
    hps = self.hps
    for group in self.param_groups:
      lr = group['lr']
      for p in group['params']:
        if p.grad is None: continue
        grad = p.grad.data
        if grad.is_sparse:
          raise RuntimeError('Shampoo does not support sparse yet')
        state = self.state[p]
        if not state:
          self.init_var_state(p, state)
        state[STEP] += 1

        preconditioner = state[PRECONDITIONER]
        graft = state[GRAFT]

        # Gather statistics, compute preconditioners
        graft.add_statistics(grad)
        if state[STEP] % hps.statistics_compute_steps == 0:
          preconditioner.add_statistics(grad)
        if state[STEP] % hps.preconditioning_compute_steps == 0:
          preconditioner.compute_preconditioners()

        # Precondition gradients
        graft_grad = graft.precondition_gradient(grad)
        shampoo_grad = grad
        if state[STEP] >= self.hps.start_preconditioning_step:
          shampoo_grad = preconditioner.preconditioned_grad(grad)

        # Grafting
        graft_norm = torch.linalg.norm(graft_grad)
        shampoo_norm = torch.linalg.norm(shampoo_grad)
        shampoo_grad.mul_(graft_norm / (shampoo_norm + 1e-16))

        # Weight decay
        if self.hps.weight_decay != 0.0:
          shampoo_grad.add_(p.data, alpha=self.hps.weight_decay)
          graft_grad.add_(p.data, alpha=self.hps.weight_decay)

        # Momentum and Nesterov momentum, if needed
        state[MOMENTUM].mul_(group['momentum']).add_(shampoo_grad)
        graft_momentum = graft.update_momentum(grad, group['momentum'])

        if state[STEP] >= self.hps.start_preconditioning_step:
          momentum_update = state[MOMENTUM]
          wd_update = shampoo_grad
        else:
          momentum_update = graft_momentum
          wd_update = graft_grad

        if hps.nesterov:
          momentum_update.mul_(group['momentum']).add_(wd_update)

        # Final update
        p.data.add_(momentum_update, alpha=-lr)
//  [python]
def _orthogonalize_gram_schmidt(matrices, epsilon=0):
    """
    Applies Gram-Schmidt procedure to orthogonalize a batch of matrices.
    If epsilon is 0, this is equivalent to `torch.qr(matrices, out=(matrices, _))`,
    """
    num_cols = matrices.shape[2]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrices[:, :, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input batch of matrices covers the gradients of at least one entire layer
        # in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.norm(col, dim=1, keepdim=True)
            except ZeroDivisionError:
                logger.error(
                    "The matrices to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.norm(col, dim=1, keepdim=True) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrices[:, :, i + 1 :]
            rest -= torch.sum(col * rest, dim=1, keepdim=True) * col
//  [python]
def _orthogonalize_gram_schmidt(matrices, epsilon=0):
    """
    Applies Gram-Schmidt procedure to orthogonalize a batch of matrices.
    If epsilon is 0, this is equivalent to `torch.qr(matrices, out=(matrices, _))`,
    """
    num_cols = matrices.shape[2]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrices[:, :, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input batch of matrices covers the gradients of at least one entire layer
        # in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.linalg.norm(col, dim=1, keepdim=True)
            except ZeroDivisionError:
                logger.error(
                    "The matrices to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.linalg.norm(col, dim=1, keepdim=True) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrices[:, :, i + 1 :]
            rest -= torch.sum(col * rest, dim=1, keepdim=True) * col
//  [python]
def forward(self, mels: torch.FloatTensor):
        _, (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):
        _, (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 forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Compute the Coulomb energy.

        Args:
            inputs (dict(str,torch.Tensor)): Input batch.

        Returns:
            dict(str, torch.Tensor): results with Coulomb energy.
        """
        q = inputs[self.charges_key].squeeze(-1)
        idx_m = inputs[properties.idx_m]

        if self.use_neighbors_lr:
            r_ij = inputs[properties.Rij_lr]
            idx_i = inputs[properties.idx_i_lr]
            idx_j = inputs[properties.idx_j_lr]
        else:
            r_ij = inputs[properties.Rij]
            idx_i = inputs[properties.idx_i]
            idx_j = inputs[properties.idx_j]

        d_ij = torch.norm(r_ij, dim=1)

        n_atoms = q.shape[0]
        n_molecules = int(idx_m[-1]) + 1

        q_ij = q[idx_i] * q[idx_j]

        potential = self.coulomb_potential(d_ij)

        # Apply cutoff if requested (shifting to zero)
        if self.cutoff is not None:
            potential = potential + self.shift**2 / potential - 2.0 * self.shift
            potential = torch.where(
                d_ij <= self.cutoff, potential, torch.zeros_like(potential)
            )

        y = snn.scatter_add((q_ij * potential), idx_i, dim_size=n_atoms)
        y = snn.scatter_add(y, idx_m, dim_size=n_molecules)
        y = 0.5 * self.ke * torch.squeeze(y, -1)

        inputs[self.output_key] = y
        return inputs
//  [python]
def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Compute the Coulomb energy.

        Args:
            inputs (dict(str,torch.Tensor)): Input batch.

        Returns:
            dict(str, torch.Tensor): results with Coulomb energy.
        """
        q = inputs[self.charges_key].squeeze(-1)
        idx_m = inputs[properties.idx_m]

        if self.use_neighbors_lr:
            r_ij = inputs[properties.Rij_lr]
            idx_i = inputs[properties.idx_i_lr]
            idx_j = inputs[properties.idx_j_lr]
        else:
            r_ij = inputs[properties.Rij]
            idx_i = inputs[properties.idx_i]
            idx_j = inputs[properties.idx_j]

        d_ij = torch.linalg.norm(r_ij, dim=1)

        n_atoms = q.shape[0]
        n_molecules = int(idx_m[-1]) + 1

        q_ij = q[idx_i] * q[idx_j]

        potential = self.coulomb_potential(d_ij)

        # Apply cutoff if requested (shifting to zero)
        if self.cutoff is not None:
            potential = potential + self.shift**2 / potential - 2.0 * self.shift
            potential = torch.where(
                d_ij <= self.cutoff, potential, torch.zeros_like(potential)
            )

        y = snn.scatter_add((q_ij * potential), idx_i, dim_size=n_atoms)
        y = snn.scatter_add(y, idx_m, dim_size=n_molecules)
        y = 0.5 * self.ke * torch.squeeze(y, -1)

        inputs[self.output_key] = y
        return inputs
//  [python]
def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Compute the Coulomb energy of the periodic system.

        Args:
            inputs (dict(str,torch.Tensor)): Input batch.

        Returns:
            dict(str, torch.Tensor): results with Coulomb energy.
        """
        q = inputs[self.charges_key].squeeze(-1)
        idx_m = inputs[properties.idx_m]

        # Use long range neighbor list if requested
        if self.use_neighbors_lr:
            r_ij = inputs[properties.Rij_lr]
            idx_i = inputs[properties.idx_i_lr]
            idx_j = inputs[properties.idx_j_lr]
        else:
            r_ij = inputs[properties.Rij]
            idx_i = inputs[properties.idx_i]
            idx_j = inputs[properties.idx_j]

        d_ij = torch.norm(r_ij, dim=1)

        positions = inputs[properties.R]
        cell = inputs[properties.cell]

        n_atoms = q.shape[0]
        n_molecules = int(idx_m[-1]) + 1

        # Get real space and reciprocal space contributions
        y_real = self._real_space(q, d_ij, idx_i, idx_j, idx_m, n_atoms, n_molecules)
        y_reciprocal = self._reciprocal_space(q, positions, cell, idx_m, n_molecules)

        y = y_real + y_reciprocal

        inputs[self.output_key] = y
        return inputs
//  [python]
def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Compute the Coulomb energy of the periodic system.

        Args:
            inputs (dict(str,torch.Tensor)): Input batch.

        Returns:
            dict(str, torch.Tensor): results with Coulomb energy.
        """
        q = inputs[self.charges_key].squeeze(-1)
        idx_m = inputs[properties.idx_m]

        # Use long range neighbor list if requested
        if self.use_neighbors_lr:
            r_ij = inputs[properties.Rij_lr]
            idx_i = inputs[properties.idx_i_lr]
            idx_j = inputs[properties.idx_j_lr]
        else:
            r_ij = inputs[properties.Rij]
            idx_i = inputs[properties.idx_i]
            idx_j = inputs[properties.idx_j]

        d_ij = torch.linalg.norm(r_ij, dim=1)

        positions = inputs[properties.R]
        cell = inputs[properties.cell]

        n_atoms = q.shape[0]
        n_molecules = int(idx_m[-1]) + 1

        # Get real space and reciprocal space contributions
        y_real = self._real_space(q, d_ij, idx_i, idx_j, idx_m, n_atoms, n_molecules)
        y_reciprocal = self._reciprocal_space(q, positions, cell, idx_m, n_molecules)

        y = y_real + y_reciprocal

        inputs[self.output_key] = y
        return inputs
//  [python]
def linear(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        # return B * R * SR * K
        if axis_weight is None or (axis_weight[..., 0] == 1 and axis_weight[..., 2] ==1) :
            weights = 1. / torch.clamp(torch.norm(dists[..., :3], dim=-1), min= 1e-6)
        else:
            weights = 1. / torch.clamp(torch.sqrt(torch.sum(torch.square(dists[...,:2]), dim=-1)) * axis_weight[..., 0] + torch.abs(dists[...,2]) * axis_weight[..., 1], min= 1e-6)
        weights = pnt_mask * weights
        return weights, embedding
//  [python]
def linear(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        # return B * R * SR * K
        if axis_weight is None or (axis_weight[..., 0] == 1 and axis_weight[..., 2] ==1) :
            weights = 1. / torch.clamp(torch.linalg.norm(dists[..., :3], dim=-1), min= 1e-6)
        else:
            weights = 1. / torch.clamp(torch.sqrt(torch.sum(torch.square(dists[...,:2]), dim=-1)) * axis_weight[..., 0] + torch.abs(dists[...,2]) * axis_weight[..., 1], min= 1e-6)
        weights = pnt_mask * weights
        return weights, embedding
//  [python]
def test_norm_fro_2_equivalence_old(self, device, dtype):
        input_sizes = [
            (0,),
            (10,),
            (0, 0),
            (4, 30),
            (0, 45),
            (100, 0),
            (45, 10, 23),
            (0, 23, 59),
            (23, 0, 37),
            (34, 58, 0),
            (0, 0, 348),
            (0, 3434, 0),
            (0, 0, 0),
            (5, 3, 8, 1, 3, 5)]

        for input_size in input_sizes:
            a = make_tensor(input_size, dtype=dtype, device=device, low=-9, high=9)

            # Try full reduction
            dim_settings = [None]

            # Try all possible 1-D reductions
            dim_settings += list(range(-a.dim(), a.dim()))

            def wrap_dim(dim, ndims):
                assert (dim < ndims) and (dim >= -ndims)
                if dim >= 0:
                    return dim
                else:
                    return dim + ndims

            # Try all possible 2-D reductions
            dim_settings += [
                (d0, d1) for d0, d1 in itertools.combinations(range(-a.dim(), a.dim()), 2)
                if wrap_dim(d0, a.dim()) != wrap_dim(d1, a.dim())]

            for dim in dim_settings:
                for keepdim in [True, False]:
                    a_norm_2 = torch.norm(a, p=2, dim=dim, keepdim=keepdim)
                    a_norm_fro = torch.norm(a, p='fro', dim=dim, keepdim=keepdim)
                    self.assertEqual(a_norm_fro, a_norm_2)
//  [python]
def test_norm_fro_2_equivalence_old(self, device, dtype):
        input_sizes = [
            (0,),
            (10,),
            (0, 0),
            (4, 30),
            (0, 45),
            (100, 0),
            (45, 10, 23),
            (0, 23, 59),
            (23, 0, 37),
            (34, 58, 0),
            (0, 0, 348),
            (0, 3434, 0),
            (0, 0, 0),
            (5, 3, 8, 1, 3, 5)]

        for input_size in input_sizes:
            a = make_tensor(input_size, dtype=dtype, device=device, low=-9, high=9)

            # Try full reduction
            dim_settings = [None]

            # Try all possible 1-D reductions
            dim_settings += list(range(-a.dim(), a.dim()))

            def wrap_dim(dim, ndims):
                assert (dim < ndims) and (dim >= -ndims)
                if dim >= 0:
                    return dim
                else:
                    return dim + ndims

            # Try all possible 2-D reductions
            dim_settings += [
                (d0, d1) for d0, d1 in itertools.combinations(range(-a.dim(), a.dim()), 2)
                if wrap_dim(d0, a.dim()) != wrap_dim(d1, a.dim())]

            for dim in dim_settings:
                for keepdim in [True, False]:
                    a_norm_2 = torch.linalg.norm(a, ord=2, dim=dim, keepdim=keepdim)
                    a_norm_fro = torch.linalg.norm(a, ord='fro', dim=dim, keepdim=keepdim)
                    self.assertEqual(a_norm_fro, a_norm_2)
//  [python]
def loss_8c(self, normal):
        """
        normal: B x N x 3
        friction: B x N x 3
        x: B x N x 3
        mu: ()
        """
        normal = normal / torch.norm(normal, dim=-1, keepdim=True)
        left = torch.einsum('ijk, ijk->ij', normal, normal)
        right = torch.norm(normal, dim=2) / self.sqrt_sq_mu_1
        diff = left - right
        return torch.sum(self.relu(-diff), 1)
//  [python]
def loss_8c(self, normal):
        """
        normal: B x N x 3
        friction: B x N x 3
        x: B x N x 3
        mu: ()
        """
        normal = normal / torch.linalg.norm(normal, dim=-1, keepdim=True)
        left = torch.einsum('ijk, ijk->ij', normal, normal)
        right = torch.linalg.norm(normal, dim=2) / self.sqrt_sq_mu_1
        diff = left - right
        return torch.sum(self.relu(-diff), 1)
//  [python]
def project(data, pgd_coef=1):
    """ project the data to a unit ball """
    data_norm = torch.reshape(torch.norm(torch.flatten(data, start_dim=1, end_dim=-1), dim=-1),
                              [data.shape[0], *[1] * (data.dim() - 1)])
    return data / data_norm * pgd_coef
//  [python]
def project(data, pgd_coef=1):
    """ project the data to a unit ball """
    data_norm = torch.reshape(torch.linalg.norm(torch.flatten(data, start_dim=1, end_dim=-1), dim=-1),
                              [data.shape[0], *[1] * (data.dim() - 1)])
    return data / data_norm * pgd_coef
//  [python]
def numlinear(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        # return B * R * SR * K
        if axis_weight is None or (axis_weight[..., 0] == 1 and axis_weight[..., 2] ==1) :
            weights = 1. / torch.clamp(torch.norm(dists, dim=-1), min= 1e-6)
        else:
            weights = 1. / torch.clamp(torch.sqrt(torch.sum(torch.square(dists[...,:2]), dim=-1)) * axis_weight[..., 0] + torch.abs(dists[...,2]) * axis_weight[..., 1], min= 1e-6)
        weights = pnt_mask * weights
        norm_weights = weights / torch.clamp(torch.sum(pnt_mask, dim=-1, keepdim=True), min=1)
        return norm_weights, embedding
//  [python]
def numlinear(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        # return B * R * SR * K
        if axis_weight is None or (axis_weight[..., 0] == 1 and axis_weight[..., 2] ==1) :
            weights = 1. / torch.clamp(torch.linalg.norm(dists, dim=-1), min= 1e-6)
        else:
            weights = 1. / torch.clamp(torch.sqrt(torch.sum(torch.square(dists[...,:2]), dim=-1)) * axis_weight[..., 0] + torch.abs(dists[...,2]) * axis_weight[..., 1], min= 1e-6)
        weights = pnt_mask * weights
        norm_weights = weights / torch.clamp(torch.sum(pnt_mask, dim=-1, keepdim=True), min=1)
        return norm_weights, embedding
//  [python]
def dist_mat(self, x):

        try:
            x = th.from_numpy(x)
        except TypeError:
            x = x

        if len(x.size()) == 4:
            x = x.view(x.size()[0], -1)

        dist = th.norm(x[:, None] - x, dim=2)
        return dist
//  [python]
def dist_mat(self, x):

        try:
            x = th.from_numpy(x)
        except TypeError:
            x = x

        if len(x.size()) == 4:
            x = x.view(x.size()[0], -1)

        dist = th.linalg.norm(x[:, None] - x, dim=2)
        return dist
//  [python]
def _greedy_diversity_order(self, query: str, documents: List[Document]) -> List[Document]:
        """
        Orders the given list of documents to maximize diversity.

        The algorithm first calculates embeddings for each document and the query. It starts by selecting the document
        that is semantically closest to the query. Then, for each remaining document, it selects the one that, on
        average, is least similar to the already selected documents. This process continues until all documents are
        selected, resulting in a list where each subsequent document contributes the most to the overall diversity of
        the selected set.

        :param query: The search query.
        :param documents: The list of Document objects to be ranked.

        :return: A list of documents ordered to maximize diversity.
        """
        texts_to_embed = self._prepare_texts_to_embed(documents)

        # Calculate embeddings
        doc_embeddings = self.model.encode(texts_to_embed, convert_to_tensor=True)  # type: ignore[attr-defined]
        query_embedding = self.model.encode([self.query_prefix + query + self.query_suffix], convert_to_tensor=True)  # type: ignore[attr-defined]

        # Normalize embeddings to unit length for computing cosine similarity
        if self.similarity == "cosine":
            doc_embeddings /= torch.norm(doc_embeddings, p=2, dim=-1).unsqueeze(-1)
            query_embedding /= torch.norm(query_embedding, p=2, dim=-1).unsqueeze(-1)

        n = len(documents)
        selected: List[int] = []

        # Compute the similarity vector between the query and documents
        query_doc_sim = query_embedding @ doc_embeddings.T

        # Start with the document with the highest similarity to the query
        selected.append(int(torch.argmax(query_doc_sim).item()))

        selected_sum = doc_embeddings[selected[0]] / n

        while len(selected) < n:
            # Compute mean of dot products of all selected documents and all other documents
            similarities = selected_sum @ doc_embeddings.T
            # Mask documents that are already selected
            similarities[selected] = torch.inf
            # Select the document with the lowest total similarity score
            index_unselected = int(torch.argmin(similarities).item())
            selected.append(index_unselected)
            # It's enough just to add to the selected vectors because dot product is distributive
            # It's divided by n for numerical stability
            selected_sum += doc_embeddings[index_unselected] / n

        ranked_docs: List[Document] = [documents[i] for i in selected]

        return ranked_docs
//  [python]
def _greedy_diversity_order(self, query: str, documents: List[Document]) -> List[Document]:
        """
        Orders the given list of documents to maximize diversity.

        The algorithm first calculates embeddings for each document and the query. It starts by selecting the document
        that is semantically closest to the query. Then, for each remaining document, it selects the one that, on
        average, is least similar to the already selected documents. This process continues until all documents are
        selected, resulting in a list where each subsequent document contributes the most to the overall diversity of
        the selected set.

        :param query: The search query.
        :param documents: The list of Document objects to be ranked.

        :return: A list of documents ordered to maximize diversity.
        """
        texts_to_embed = self._prepare_texts_to_embed(documents)

        # Calculate embeddings
        doc_embeddings = self.model.encode(texts_to_embed, convert_to_tensor=True)  # type: ignore[attr-defined]
        query_embedding = self.model.encode([self.query_prefix + query + self.query_suffix], convert_to_tensor=True)  # type: ignore[attr-defined]

        # Normalize embeddings to unit length for computing cosine similarity
        if self.similarity == "cosine":
            doc_embeddings /= torch.linalg.norm(doc_embeddings, p=2, dim=-1, keepdim=True)
            query_embedding /= torch.linalg.norm(query_embedding, p=2, dim=-1, keepdim=True)

        n = len(documents)
        selected: List[int] = []

        # Compute the similarity vector between the query and documents
        query_doc_sim = query_embedding @ doc_embeddings.T

        # Start with the document with the highest similarity to the query
        selected.append(int(torch.argmax(query_doc_sim).item()))

        selected_sum = doc_embeddings[selected[0]] / n

        while len(selected) < n:
            # Compute mean of dot products of all selected documents and all other documents
            similarities = selected_sum @ doc_embeddings.T
            # Mask documents that are already selected
            similarities[selected] = torch.inf
            # Select the document with the lowest total similarity score
            index_unselected = int(torch.argmin(similarities).item())
            selected.append(index_unselected)
            # It's enough just to add to the selected vectors because dot product is distributive
            # It's divided by n for numerical stability
            selected_sum += doc_embeddings[index_unselected] / n

        ranked_docs: List[Document] = [documents[i] for i in selected]

        return ranked_docs
//  [python]
def kernel_loss(self, k_x, k_y, k_l):

        beta = 1.0

        L = th.norm(k_l)
        Y = th.norm(k_y) ** beta
        X = th.norm(k_x) ** (1-beta)

        LY = th.trace(th.matmul(k_l, k_y))**beta
        LX = th.trace(th.matmul(k_l, k_x))**(1-beta)

        return 2*th.log2((LY*LX)/(L*Y*X))
//  [python]
def kernel_loss(self, k_x, k_y, k_l):

        beta = 1.0

        L = th.linalg.norm(k_l)
        Y = th.linalg.norm(k_y) ** beta
        X = th.linalg.norm(k_x) ** (1-beta)

        LY = th.trace(th.matmul(k_l, k_y))**beta
        LX = th.trace(th.matmul(k_l, k_x))**(1-beta)

        return 2*th.log2((LY*LX)/(L*Y*X))
//  [python]
def distance_correlation(A, B):
    return 1 - torch.bmm(A, B).diagonal(0, -1, -2).sum(-1) / (torch.norm(A, dim=(1, 2)) * torch.norm(B, dim=(1, 2)))
//  [python]
def distance_correlation(A, B):
    return 1 - torch.bmm(A, B).diagonal(0, -1, -2).sum(-1) / (torch.linalg.norm(A, dim=(1, 2)) * torch.linalg.norm(B, dim=(1, 2)))
//  [python]
def centerDatas(datas):
    datas= datas - datas.mean(1, keepdim=True)
    datas = datas / torch.norm(datas, dim=2, keepdim= True)
    # datas[:, n_lsamples:] = datas[:, n_lsamples:, :] - datas[:, n_lsamples:].mean(1, keepdim=True)
    # datas[:, n_lsamples:] = datas[:, n_lsamples:, :] / torch.norm(datas[:, n_lsamples:, :], 2, 2)[:, :, None]

    return datas
//  [python]
def centerDatas(datas):
    datas= datas - datas.mean(1, keepdim=True)
    datas = datas / torch.linalg.norm(datas, dim=2, keepdim= True)
    # datas[:, n_lsamples:] = datas[:, n_lsamples:, :] - datas[:, n_lsamples:].mean(1, keepdim=True)
    # datas[:, n_lsamples:] = datas[:, n_lsamples:, :] / torch.linalg.norm(datas[:, n_lsamples:, :], 2, 2)[:, :, None]

    return datas
//  [python]
def test_PCACompute(self):
        vec1 = torch.randn(3)
        vec1 = vec1 / torch.norm(vec1)
        vec2 = torch.randn(3)
        vec2 = vec2 / torch.norm(vec2)
        norm = vec1.cross(vec2) / torch.norm(vec1.cross(vec2))
        plane = torch.randn(100, 1) * vec1.view(1, 3) + torch.randn(100, 1) * vec2.view(1, 3)
        data = Data(pos=plane)
        pca = PCACompute()
        data = pca(data)
        npt.assert_allclose(np.abs(data.eigenvectors[:, 0].dot(norm).item()), 1.0, atol=1e-5)
//  [python]
def test_PCACompute(self):
        vec1 = torch.randn(3)
        vec1 = vec1 / torch.linalg.norm(vec1)
        vec2 = torch.randn(3)
        vec2 = vec2 / torch.linalg.norm(vec2)
        norm = vec1.cross(vec2) / torch.linalg.norm(vec1.cross(vec2))
        plane = torch.randn(100, 1) * vec1.view(1, 3) + torch.randn(100, 1) * vec2.view(1, 3)
        data = Data(pos=plane)
        pca = PCACompute()
        data = pca(data)
        npt.assert_allclose(np.abs(data.eigenvectors[:, 0].dot(norm).item()), 1.0, atol=1e-5)
//  [python]
def norm(x, ord='fro', axis=None, keepdims=False):
    return torch.norm(x, p=ord, dim=axis, keepdim=keepdims)
//  [python]
def norm(x, ord='fro', axis=None, keepdims=False):
    return torch.linalg.norm(x, ord=ord, dim=axis, keepdim=keepdims)
//  [python]
def update_via_matrix(R, q, update_mask=None):
    if update_mask is None: update_mask = torch.ones_like(q[..., 0])

    # Compute normalized quaternion
    q_normalized = torch.ones(q.shape[:-1] + (4,), device=q.device, dtype=q.dtype)
    q_normalized[..., 1:] = q*update_mask
    q_normalized = q_normalized / torch.norm(q_normalized, dim=-1, keepdim=True)

    # Compute angle associated with rotation
    angle = 2 * torch.acos(q_normalized[..., 0])
    axis = q/torch.norm(q, dim=-1, keepdim=True)
    
    # Udate rotation matrix via right multiplication of delta rotation
    rot_vec = angle[..., None] * axis
    delta_R = so3_utils.Exp(rot_vec)
    return torch.matmul(R, delta_R)
//  [python]
def update_via_matrix(R, q, update_mask=None):
    if update_mask is None: update_mask = torch.ones_like(q[..., 0])

    # Compute normalized quaternion
    q_normalized = torch.ones(q.shape[:-1] + (4,), device=q.device, dtype=q.dtype)
    q_normalized[..., 1:] = q*update_mask
    q_normalized = q_normalized / torch.linalg.norm(q_normalized, dim=-1, keepdim=True)

    # Compute angle associated with rotation
    angle = 2 * torch.acos(q_normalized[..., 0])
    axis = q/torch.linalg.norm(q, dim=-1, keepdim=True)
    
    # Udate rotation matrix via right multiplication of delta rotation
    rot_vec = angle[..., None] * axis
    delta_R = so3_utils.Exp(rot_vec)
    return torch.matmul(R, delta_R)
//  [python]
def expand(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        init_pos = batched_data["init_pos"] # B x T x 3
        atoms = batched_data["x"][:, :, 0]
        batch_size, max_num_atoms = pos.size()[:2]
        cell = batched_data["cell"] # B x 3 x 3
        cell_tensor = torch.tensor(self.cells, device=pos.device).to(cell.dtype).unsqueeze(0).expand(batch_size, -1, -1)
        offset = torch.bmm(cell_tensor, cell) # B x 8 x 3
        expand_pos = pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        expand_pos = expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        init_expand_pos = init_pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        init_expand_pos = init_expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        expand_dist = torch.norm(pos.unsqueeze(2) - expand_pos.unsqueeze(1), p=2, dim=-1) # B x T x (8 x T)
        expand_mask = expand_dist < self.cutoff # B x T x (8 x T)
        expand_mask = torch.masked_fill(expand_mask, atoms.eq(0).unsqueeze(-1), False)
        expand_mask = (torch.sum(expand_mask, dim=1) > 0) & (~(atoms.eq(0).repeat(1, 8))) # B x (8 x T)
        expand_len = torch.sum(expand_mask, dim=-1)
        max_expand_len = torch.max(expand_len)
        outcell_index = torch.zeros([batch_size, max_expand_len], dtype=torch.long, device=pos.device)
        expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        init_expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        outcell_all_index = torch.arange(max_num_atoms, dtype=torch.long, device=pos.device).repeat(len(self.cells))
        for i in range(batch_size):
           outcell_index[i, :expand_len[i]] = outcell_all_index[expand_mask[i]]
           expand_pos_compressed[i, :expand_len[i], :] = expand_pos[i, expand_mask[i], :]
           init_expand_pos_compressed[i, :expand_len[i], :] = init_expand_pos[i, expand_mask[i], :]
        batched_data["expand_pos"] = expand_pos_compressed
        batched_data["init_expand_pos"] = init_expand_pos_compressed
        batched_data["expand_len"] = expand_len
        batched_data["outcell_index"] = outcell_index
        batched_data["expand_mask"] = mask_after_k_persample(batch_size, max_expand_len, expand_len)

        return batched_data
//  [python]
def expand(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        init_pos = batched_data["init_pos"] # B x T x 3
        atoms = batched_data["x"][:, :, 0]
        batch_size, max_num_atoms = pos.size()[:2]
        cell = batched_data["cell"] # B x 3 x 3
        cell_tensor = torch.tensor(self.cells, device=pos.device).to(cell.dtype).unsqueeze(0).expand(batch_size, -1, -1)
        offset = torch.bmm(cell_tensor, cell) # B x 8 x 3
        expand_pos = pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        expand_pos = expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        init_expand_pos = init_pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        init_expand_pos = init_expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        expand_dist = torch.linalg.norm(pos.unsqueeze(2) - expand_pos.unsqueeze(1), ord=2, dim=-1) # B x T x (8 x T)
        expand_mask = expand_dist < self.cutoff # B x T x (8 x T)
        expand_mask = torch.masked_fill(expand_mask, atoms.eq(0).unsqueeze(-1), False)
        expand_mask = (torch.sum(expand_mask, dim=1) > 0) & (~(atoms.eq(0).repeat(1, 8))) # B x (8 x T)
        expand_len = torch.sum(expand_mask, dim=-1)
        max_expand_len = torch.max(expand_len)
        outcell_index = torch.zeros([batch_size, max_expand_len], dtype=torch.long, device=pos.device)
        expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        init_expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        outcell_all_index = torch.arange(max_num_atoms, dtype=torch.long, device=pos.device).repeat(len(self.cells))
        for i in range(batch_size):
           outcell_index[i, :expand_len[i]] = outcell_all_index[expand_mask[i]]
           expand_pos_compressed[i, :expand_len[i], :] = expand_pos[i, expand_mask[i], :]
           init_expand_pos_compressed[i, :expand_len[i], :] = init_expand_pos[i, expand_mask[i], :]
        batched_data["expand_pos"] = expand_pos_compressed
        batched_data["init_expand_pos"] = init_expand_pos_compressed
        batched_data["expand_len"] = expand_len
        batched_data["outcell_index"] = outcell_index
        batched_data["expand_mask"] = mask_after_k_persample(batch_size, max_expand_len, expand_len)

        return batched_data
//  [python]
def mean_curvature(y, x):
    grad = gradient(y, x)
    grad_norm = torch.norm(grad, dim=-1)
    unit_grad = grad.squeeze(-1)/grad_norm.unsqueeze(-1)

    Km = 0.5*divergence(unit_grad, x)
    return Km
//  [python]
def mean_curvature(y, x):
    grad = gradient(y, x)
    grad_norm = torch.linalg.norm(grad, dim=-1)
    unit_grad = grad.squeeze(-1)/grad_norm.unsqueeze(-1)

    Km = 0.5*divergence(unit_grad, x)
    return Km
//  [python]
def compute_jitter(pred_output, fps=30):
    """compute jitter of the motion
    Args:
        pred_output (SMPL ModelOutput).
        fps (float).
    Returns:
        jitter (N-3).
    """
    
    pred3d = pred_output.joints[:, :24]
        
    pred_jitter = torch.norm(
        (pred3d[3:] - 3 * pred3d[2:-1] + 3 * pred3d[1:-2] - pred3d[:-3]) * (fps**3),
        dim=2,
    ).mean(dim=-1)
    
    return pred_jitter.cpu().numpy() / 10.0
//  [python]
def compute_jitter(pred_output, fps=30):
    """compute jitter of the motion
    Args:
        pred_output (SMPL ModelOutput).
        fps (float).
    Returns:
        jitter (N-3).
    """
    
    pred3d = pred_output.joints[:, :24]
        
    pred_jitter = torch.linalg.norm(
        (pred3d[3:] - 3 * pred3d[2:-1] + 3 * pred3d[1:-2] - pred3d[:-3]) * (fps**3),
        dim=2,
    ).mean(dim=-1)
    
    return pred_jitter.cpu().numpy() / 10.0
//  [python]
def compute_rte(target_trans, pred_trans):
    # Compute the global alignment
    _, rot, trans = align_pcl(target_trans[None, :], pred_trans[None, :], fixed_scale=True)
    pred_trans_hat = (
        torch.einsum("tij,tnj->tni", rot, pred_trans[None, :]) + trans[None, :]
    )[0]
    
    # Compute the entire displacement of ground truth trajectory
    disps, disp = [], 0
    for p1, p2 in zip(target_trans, target_trans[1:]):
        delta = (p2 - p1).norm(2, dim=-1)
        disp += delta
        disps.append(disp)
    
    # Compute absolute root-translation-error (RTE)
    rte = torch.norm(target_trans - pred_trans_hat, 2, dim=-1)
    
    # Normalize it to the displacement
    return (rte / disp).numpy()
//  [python]
def compute_rte(target_trans, pred_trans):
    # Compute the global alignment
    _, rot, trans = align_pcl(target_trans[None, :], pred_trans[None, :], fixed_scale=True)
    pred_trans_hat = (
        torch.einsum("tij,tnj->tni", rot, pred_trans[None, :]) + trans[None, :]
    )[0]
    
    # Compute the entire displacement of ground truth trajectory
    disps, disp = [], 0
    for p1, p2 in zip(target_trans, target_trans[1:]):
        delta = torch.linalg.norm(p2 - p1, 2, dim=-1)
        disp += delta
        disps.append(disp)
    
    # Compute absolute root-translation-error (RTE)
    rte = torch.linalg.norm(target_trans - pred_trans_hat, 2, dim=-1)
    
    # Normalize it to the displacement
    return (rte / disp).numpy()
//  [python]
def compute_weights(pts0, pts1, par):
    return par / (torch.norm(pts0 - pts1, dim=1).unsqueeze(1) + par)
//  [python]
def compute_weights(pts0, pts1, par):
    return par / (torch.linalg.norm(pts0 - pts1, dim=1).unsqueeze(1) + par)
//  [python]
def tensor_curvature(y, x):
    grad = gradient(y, x)
    grad_norm = torch.norm(grad, dim=-1)
    unit_grad = grad.squeeze(-1)/grad_norm.unsqueeze(-1)

    T = -jacobian(unit_grad, x)[0]

    print(T)
    e, v = torch.eig(T, eigenvectors=True)

    print(e)
    print(v)

    return T
//  [python]
def tensor_curvature(y, x):
    grad = gradient(y, x)
    grad_norm = torch.linalg.norm(grad, dim=-1)
    unit_grad = grad.squeeze(-1)/grad_norm.unsqueeze(-1)

    T = -jacobian(unit_grad, x)[0]

    print(T)
    e, v = torch.eig(T, eigenvectors=True)

    print(e)
    print(v)

    return T
//  [python]
def __compute_distance(self, a, b, order):
        if order == 'inf':
            order = np.inf
        if order == 'cos' or order == 'cosine':
            return 1.0 - torch.nn.CosineSimilarity(dim=0)(a, b).cpu().numpy()
        else:
            if order != np.inf:
                try:
                    order = int(order)
                except:
                    raise TypeError("input metric for distance is not understandable")
            return torch.norm(a - b, p=order).cpu().numpy()
//  [python]
def __compute_distance(self, a, b, order):
        if order == 'inf':
            order = np.inf
        if order == 'cos' or order == 'cosine':
            return 1.0 - torch.nn.CosineSimilarity(dim=0)(a, b).cpu().numpy()
        else:
            if order != np.inf:
                try:
                    order = int(order)
                except:
                    raise TypeError("input metric for distance is not understandable")
            return torch.linalg.norm(a - b, ord=order).cpu().numpy()