jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def inn_pred(W, X, A, B, phi, mitr=300, tol=3e-6, trasposed_A=False, compute_dphi=False, name='Forward'):
        # TODO: randomized speed up
        At = A if trasposed_A else torch.transpose(A, 0, 1)
        #X = B if X is None else X
        results = dict()
        results['abs_err'] = []
        results['X_new'] = []
        results['X'] = []
        results['X_diff'] = []
        err = 0
        status = 'max itrs reached'
        for i in range(mitr):
            # WXA
            X_ = W @ X
            support = torch.spmm(At, X_.T).T
            # tmp = support + B
            # X_new = phi(F.normalize(tmp, dim=0))
            X_new = phi(support + B)
            err = torch.norm(X_new - X, np.inf)
            results['abs_err'].append(err.item())
            results['X_diff'].append((X_new - X).cpu().numpy())
            results['X_new'].append(X_new.cpu().numpy())
            results['X'].append(X.cpu().numpy())
            if err < tol:
                status = 'converged'
                break
            X = X_new
        # print(f'{name}, no. iterations: {i}, err: {err}')
        dphi = None
        if compute_dphi:
            with torch.enable_grad():
                support = torch.spmm(At, (W @ X).T).T
                Z = support + B
                Z.requires_grad_(True)
                X_new = phi(Z)
                dphi = torch.autograd.grad(torch.sum(X_new), Z, only_inputs=True)[0]

        return X_new, err, status, dphi, results
//  [python]
def inn_pred(W, X, A, B, phi, mitr=300, tol=3e-6, trasposed_A=False, compute_dphi=False, name='Forward'):
        # TODO: randomized speed up
        At = A if trasposed_A else torch.transpose(A, 0, 1)
        #X = B if X is None else X
        results = dict()
        results['abs_err'] = []
        results['X_new'] = []
        results['X'] = []
        results['X_diff'] = []
        err = 0
        status = 'max itrs reached'
        for i in range(mitr):
            # WXA
            X_ = W @ X
            support = torch.spmm(At, X_.T).T
            # tmp = support + B
            # X_new = phi(F.normalize(tmp, dim=0))
            X_new = phi(support + B)
            err = torch.linalg.norm(X_new - X, np.inf)
            results['abs_err'].append(err.item())
            results['X_diff'].append((X_new - X).cpu().numpy())
            results['X_new'].append(X_new.cpu().numpy())
            results['X'].append(X.cpu().numpy())
            if err < tol:
                status = 'converged'
                break
            X = X_new
        # print(f'{name}, no. iterations: {i}, err: {err}')
        dphi = None
        if compute_dphi:
            with torch.enable_grad():
                support = torch.spmm(At, (W @ X).T).T
                Z = support + B
                Z.requires_grad_(True)
                X_new = phi(Z)
                dphi = torch.autograd.grad(torch.sum(X_new), Z, only_inputs=True)[0]

        return X_new, err, status, dphi, results
//  [python]
def backward(ctx, grad_output):
        # TODO
        grad_Z = grad_output
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.norm(FF, p='fro')
        # R = G * ((Q_F.t() @ B.t()) @ (grad_Y_hat @ Q_S))
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        # R = G * (Q_F.t() @ grad_Z @ torch.inverse(Q_S).t()) # revised version for asymmetric adj.
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        # ipdb.set_trace()
        # grad_X = torch.autograd.grad(Z, X, grad_outputs=grad_Z, only_inputs=True)
        # grad_X = grad_Z @ kronecker(Q_S, Q_F) @ torch.diag(G.view(-1)) @ kronecker(Q_S.t(), Q_F.t())
        # need to chagne to vectorization form, otherwise kronecker product will cause OOM.
        # grad_X = Q_F @ (G * (Q_F.t() @ Q_S)) @ Q_S.t()
        # print(f'grad_X.shape: {grad_X.shape}')
        grad_X = None
        return grad_X, grad_F, None, None, None, None
//  [python]
def backward(ctx, grad_output):
        # TODO
        grad_Z = grad_output
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.linalg.norm(FF, ord='fro')
        # R = G * ((Q_F.t() @ B.t()) @ (grad_Y_hat @ Q_S))
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        # R = G * (Q_F.t() @ grad_Z @ torch.inverse(Q_S).t()) # revised version for asymmetric adj.
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        # ipdb.set_trace()
        # grad_X = torch.autograd.grad(Z, X, grad_outputs=grad_Z, only_inputs=True)
        # grad_X = grad_Z @ kronecker(Q_S, Q_F) @ torch.diag(G.view(-1)) @ kronecker(Q_S.t(), Q_F.t())
        # need to chagne to vectorization form, otherwise kronecker product will cause OOM.
        # grad_X = Q_F @ (G * (Q_F.t() @ Q_S)) @ Q_S.t()
        # print(f'grad_X.shape: {grad_X.shape}')
        grad_X = None
        return grad_X, grad_F, None, None, None, None
//  [python]
def backward(ctx, grad_output):
        # TODO
        grad_Z = grad_output
        # ipdb.set_trace()
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.norm(FF, p='fro')
        # R = G * ((Q_F.t() @ B.t()) @ (grad_Y_hat @ Q_S))
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        grad_X = None
        return grad_X, grad_F, None, None, None, None, None, None, None
//  [python]
def backward(ctx, grad_output):
        # TODO
        grad_Z = grad_output
        # ipdb.set_trace()
        F, S, Q_F, Q_S, Z, G, X, gamma = ctx.saved_tensors
        FF = F.t() @ F
        FF_norm = torch.linalg.norm(FF, ord='fro')
        # R = G * ((Q_F.t() @ B.t()) @ (grad_Y_hat @ Q_S))
        R = G * (Q_F.t() @ grad_Z @ Q_S)
        R = Q_F @ R @ Q_S.t() @ torch.sparse.mm(S, Z.t())
        scalar_1 = gamma * (1/(FF_norm+epsilon_F))
        scalar_2 = torch.sum(FF * R)
        scalar_2 = 2 * scalar_2 * (1/(FF_norm**2 + epsilon_F * FF_norm))
        grad_F = (R + R.t()) - scalar_2 * FF
        grad_F = scalar_1 * (F @ grad_F)
        grad_X = None
        return grad_X, grad_F, None, None, None, None, None, None, None
//  [python]
def chamfer_loss(self, x, recon_x):
        with torch.no_grad():
            dist = torch.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        return torch.min(dist, dim=2).values.sum(1) + torch.min(dist, dim=1).values.sum(1)
//  [python]
def chamfer_loss(self, x, recon_x):
        with torch.no_grad():
            dist = torch.linalg.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        return torch.min(dist, dim=2).values.sum(1) + torch.min(dist, dim=1).values.sum(1)
//  [python]
def evaluation_chamfer_loss(self, x):
        with torch.no_grad():
            recon_x = self(x)
            dist = torch.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        try:
            return torch.maximum(torch.min(dist, dim=2).values.mean(1), torch.min(dist, dim=1).values.mean(1))
        except:
            cand1 = torch.min(dist, dim=2).values.mean(1)
            cand2 = torch.min(dist, dim=1).values.mean(1)
            return torch.cat([
                cand1[cand1 >= cand2], 
                cand2[cand2 > cand1]], dim=0)
//  [python]
def evaluation_chamfer_loss(self, x):
        with torch.no_grad():
            recon_x = self(x)
            dist = torch.linalg.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        try:
            return torch.maximum(torch.min(dist, dim=2).values.mean(1), torch.min(dist, dim=1).values.mean(1))
        except:
            cand1 = torch.min(dist, dim=2).values.mean(1)
            cand2 = torch.min(dist, dim=1).values.mean(1)
            return torch.cat([
                cand1[cand1 >= cand2], 
                cand2[cand2 > cand1]], dim=0)
//  [python]
def modified_chamfer_loss(self, x, recon_x):
        with torch.no_grad():
            dist = torch.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        try:
            return torch.maximum(torch.min(dist, dim=2).values.mean(1), torch.min(dist, dim=1).values.mean(1))
        except:
            cand1 = torch.min(dist, dim=2).values.mean(1)
            cand2 = torch.min(dist, dim=1).values.mean(1)
            return torch.cat([
                cand1[cand1 >= cand2], 
                cand2[cand2 > cand1]], dim=0)
//  [python]
def modified_chamfer_loss(self, x, recon_x):
        with torch.no_grad():
            dist = torch.linalg.norm(x.unsqueeze(2)-recon_x.unsqueeze(3), dim=1)
        try:
            return torch.maximum(torch.min(dist, dim=2).values.mean(1), torch.min(dist, dim=1).values.mean(1))
        except:
            cand1 = torch.min(dist, dim=2).values.mean(1)
            cand2 = torch.min(dist, dim=1).values.mean(1)
            return torch.cat([
                cand1[cand1 >= cand2], 
                cand2[cand2 > cand1]], dim=0)
//  [python]
def cal_normal(group_xyz, random_inv=False, is_group=False):
    """
    Calculate Normal Vector (Unit Form + First Term Positive)

    :param group_xyz: [B, N, K=3, 3] / [B, N, G, K=3, 3]
    :param random_inv:
    :param return_intersect:
    :param return_const:
    :return: [B, N, 3]
    """
    edge_vec1 = group_xyz[..., 1, :] - group_xyz[..., 0, :]  # [B, N, 3]
    edge_vec2 = group_xyz[..., 2, :] - group_xyz[..., 0, :]  # [B, N, 3]

    nor = torch.cross(edge_vec1, edge_vec2, dim=-1)
    unit_nor = nor / torch.norm(nor, dim=-1, keepdim=True)  # [B, N, 3] / [B, N, G, 3]
    if not is_group:
        pos_mask = (unit_nor[..., 0] > 0).float() * 2. - 1.  # keep x_n positive
    else:
        pos_mask = (unit_nor[..., 0:1, 0] > 0).float() * 2. - 1.
    unit_nor = unit_nor * pos_mask.unsqueeze(-1)

    # batch-wise random inverse normal vector (prob: 0.5)
    if random_inv:
        random_mask = torch.randint(0, 2, (group_xyz.size(0), 1, 1)).float() * 2. - 1.
        random_mask = random_mask.to(unit_nor.device)
        if not is_group:
            unit_nor = unit_nor * random_mask
        else:
            unit_nor = unit_nor * random_mask.unsqueeze(-1)

    return unit_nor
//  [python]
def cal_normal(group_xyz, random_inv=False, is_group=False):
    """
    Calculate Normal Vector (Unit Form + First Term Positive)

    :param group_xyz: [B, N, K=3, 3] / [B, N, G, K=3, 3]
    :param random_inv:
    :param return_intersect:
    :param return_const:
    :return: [B, N, 3]
    """
    edge_vec1 = group_xyz[..., 1, :] - group_xyz[..., 0, :]  # [B, N, 3]
    edge_vec2 = group_xyz[..., 2, :] - group_xyz[..., 0, :]  # [B, N, 3]

    nor = torch.cross(edge_vec1, edge_vec2, dim=-1)
    unit_nor = nor / torch.linalg.norm(nor, dim=-1, keepdim=True)  # [B, N, 3] / [B, N, G, 3]
    if not is_group:
        pos_mask = (unit_nor[..., 0] > 0).float() * 2. - 1.  # keep x_n positive
    else:
        pos_mask = (unit_nor[..., 0:1, 0] > 0).float() * 2. - 1.
    unit_nor = unit_nor * pos_mask.unsqueeze(-1)

    # batch-wise random inverse normal vector (prob: 0.5)
    if random_inv:
        random_mask = torch.randint(0, 2, (group_xyz.size(0), 1, 1)).float() * 2. - 1.
        random_mask = random_mask.to(unit_nor.device)
        if not is_group:
            unit_nor = unit_nor * random_mask
        else:
            unit_nor = unit_nor * random_mask.unsqueeze(-1)

    return unit_nor
//  [python]
def forward(self, y_hat, y):
        """Returns magnitude loss and spectral convergence loss

        Arguments
        ---------
        y_hat : torch.tensor
            generated waveform tensor
        y : torch.tensor
            real waveform tensor
        """
        y_hat_M = stft(y_hat, self.n_fft, self.hop_length, self.win_length)
        y_M = stft(y, self.n_fft, self.hop_length, self.win_length)
        loss_mag = F.l1_loss(torch.log(y_M), torch.log(y_hat_M))
        loss_sc = torch.norm(y_M - y_hat_M, p='fro') / torch.norm(y_M, p='fro')
        return loss_mag, loss_sc
//  [python]
def forward(self, y_hat, y):
        """Returns magnitude loss and spectral convergence loss

        Arguments
        ---------
        y_hat : torch.tensor
            generated waveform tensor
        y : torch.tensor
            real waveform tensor
        """
        y_hat_M = stft(y_hat, self.n_fft, self.hop_length, self.win_length)
        y_M = stft(y, self.n_fft, self.hop_length, self.win_length)
        loss_mag = F.l1_loss(torch.log(y_M), torch.log(y_hat_M))
        loss_sc = torch.linalg.norm(y_M - y_hat_M, ord='fro') / torch.linalg.norm(y_M, ord='fro')
        return loss_mag, loss_sc
//  [python]
def setUp(self):
        class TinyModel(torch.nn.Module):
            def __init__(self):
                super(TinyModel, self).__init__()
                self.linear = torch.nn.Linear(5, 10)

            def forward(self, x):
                x = self.linear(x)
                return torch.norm(x)

        self._model_class = TinyModel
//  [python]
def setUp(self):
        class TinyModel(torch.nn.Module):
            def __init__(self):
                super(TinyModel, self).__init__()
                self.linear = torch.nn.Linear(5, 10)

            def forward(self, x):
                x = self.linear(x)
                return torch.linalg.norm(x)

        self._model_class = TinyModel
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    pytroch version
    
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part last, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [axis_angle * sin_half_angles_over_angles, torch.cos(half_angles)], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle):
    """
    pytroch version
    
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part last, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = 0.5 * angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    quaternions = torch.cat(
        [axis_angle * sin_half_angles_over_angles, torch.cos(half_angles)], dim=-1
    )
    return quaternions
//  [python]
def optimizer_based_percep2d(self, init_kps3d, mview_kps2d_tensor,
                                 selected_view) -> torch.Tensor:

        optim_kps3d = nn.Parameter(init_kps3d)
        optimizer = optim.Adam([optim_kps3d], lr=self.lr)
        last_loss = torch.tensor(0.)

        for step in range(self.steps):
            optimizer.zero_grad()
            loss = torch.tensor(0.)
            proj_kps2d = self.projector.project_single_point(optim_kps3d)
            for i, (is_selected, kps2d) in enumerate(
                    zip(selected_view, mview_kps2d_tensor)):
                if is_selected and not torch.isnan(kps2d).any():
                    loss += torch.norm(proj_kps2d[i] - kps2d)

            if torch.abs(last_loss - loss) < 1e-3:
                break
            else:
                last_loss = loss
            loss.backward()
            optimizer.step()
        return optim_kps3d.detach()
//  [python]
def optimizer_based_percep2d(self, init_kps3d, mview_kps2d_tensor,
                                 selected_view) -> torch.Tensor:

        optim_kps3d = nn.Parameter(init_kps3d)
        optimizer = optim.Adam([optim_kps3d], lr=self.lr)
        last_loss = torch.tensor(0.)

        for step in range(self.steps):
            optimizer.zero_grad()
            loss = torch.tensor(0.)
            proj_kps2d = self.projector.project_single_point(optim_kps3d)
            for i, (is_selected, kps2d) in enumerate(
                    zip(selected_view, mview_kps2d_tensor)):
                if is_selected and not torch.isnan(kps2d).any():
                    loss += torch.linalg.norm(proj_kps2d[i] - kps2d)

            if torch.abs(last_loss - loss) < 1e-3:
                break
            else:
                last_loss = loss
            loss.backward()
            optimizer.step()
        return optim_kps3d.detach()
//  [python]
def fm_loss(self, z):
        bs = z.size(0)
        z_dim = z.size(1)

        # augment
        z_permuted = z[torch.randperm(bs)]
        alpha = (torch.rand(bs) * (1 + 2*self.alpha_0) - self.alpha_0).unsqueeze(1).to(z)
        z_augmented = alpha*z + (1-alpha)*z_permuted
        
        # loss
        if not self.relaxation:
            if self.use_identity:
                G = self.get_Identity_proj_Riemannian_metric(z_augmented, create_graph=True)
            else:                
                G = self.get_Fisher_proj_Riemannian_metric(z_augmented, sigma=self.sigma, create_graph=True)
            tr_G = torch.einsum('nii->n', G)
            c_squared = (tr_G/z_dim).mean()
            diff = G - c_squared * torch.eye(z_dim).to(z).unsqueeze(0)
            fm_loss = torch.norm(diff.view(bs, -1), dim=1)**2
        elif self.relaxation:
            v = torch.randn(bs, z_dim).to(z_augmented)
            X, Jv = (
                torch.autograd.functional.jvp(self.decode, z_augmented, v=v, create_graph=True)
            ) # bs num_pts 3
            if self.use_identity:
                Jv_sq_norm = torch.einsum('nij,nij->n', Jv, Jv)
            else:
                Jv_sq_norm = Info_Geo_Riem_metric(X, J=None, Xdot=Jv, sigma=self.sigma)
            TrG = Jv_sq_norm.mean()
            fm_loss = torch.mean((Jv_sq_norm - (torch.sum(v**2, dim=1)) * TrG/z_dim)**2)
        return fm_loss.sum()
//  [python]
def fm_loss(self, z):
        bs = z.size(0)
        z_dim = z.size(1)

        # augment
        z_permuted = z[torch.randperm(bs)]
        alpha = (torch.rand(bs) * (1 + 2*self.alpha_0) - self.alpha_0).unsqueeze(1).to(z)
        z_augmented = alpha*z + (1-alpha)*z_permuted
        
        # loss
        if not self.relaxation:
            if self.use_identity:
                G = self.get_Identity_proj_Riemannian_metric(z_augmented, create_graph=True)
            else:                
                G = self.get_Fisher_proj_Riemannian_metric(z_augmented, sigma=self.sigma, create_graph=True)
            tr_G = torch.einsum('nii->n', G)
            c_squared = (tr_G/z_dim).mean()
            diff = G - c_squared * torch.eye(z_dim).to(z).unsqueeze(0)
            fm_loss = torch.linalg.norm(diff.view(bs, -1), dim=1)**2
        elif self.relaxation:
            v = torch.randn(bs, z_dim).to(z_augmented)
            X, Jv = (
                torch.autograd.functional.jvp(self.decode, z_augmented, v=v, create_graph=True)
            ) # bs num_pts 3
            if self.use_identity:
                Jv_sq_norm = torch.einsum('nij,nij->n', Jv, Jv)
            else:
                Jv_sq_norm = Info_Geo_Riem_metric(X, J=None, Xdot=Jv, sigma=self.sigma)
            TrG = Jv_sq_norm.mean()
            fm_loss = torch.mean((Jv_sq_norm - (torch.sum(v**2, dim=1)) * TrG/z_dim)**2)
        return fm_loss.sum()
//  [python]
def forward(self, input):
        w = self.weight / (torch.norm(self.weight) + self.eps) * self.scale
        return F.conv2d(input, w, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)
//  [python]
def forward(self, input):
        w = self.weight / (torch.linalg.norm(self.weight) + self.eps) * self.scale
        return F.conv2d(input, w, self.bias, self.stride,
                        self.padding, self.dilation, self.groups)
//  [python]
def gromov_wasserstein(
    C1, C2, p=None, q=None, eps=0.1, max_iter=100, tol=1e-06, max_sinkhorn_iter=10
):
    m, n = C1.shape[0], C2.shape[0]
    if p is None:
        p = C1.new_ones((m,)) / m
    if q is None:
        q = C2.new_ones((n,)) / n
    T = torch.outer(p, q)

    constC, hC1, hC2 = init_matrix(C1, C2, p, q)

    cpt = 0
    err = 1.0

    while err > tol and cpt < max_iter:
        T_old = T

        # compute the gradient
        tens = gwggrad(constC, hC1, hC2, T)

        T = sinkhorn(tens, a=p, b=q, eps=eps, max_iter=max_sinkhorn_iter)

        if cpt % 10 == 0:
            # we can speed up the process by checking for the error only all
            # the 10th iterations
            err = torch.norm(T - T_old)

        cpt += 1

    return T
//  [python]
def gromov_wasserstein(
    C1, C2, p=None, q=None, eps=0.1, max_iter=100, tol=1e-06, max_sinkhorn_iter=10
):
    m, n = C1.shape[0], C2.shape[0]
    if p is None:
        p = C1.new_ones((m,)) / m
    if q is None:
        q = C2.new_ones((n,)) / n
    T = torch.outer(p, q)

    constC, hC1, hC2 = init_matrix(C1, C2, p, q)

    cpt = 0
    err = 1.0

    while err > tol and cpt < max_iter:
        T_old = T

        # compute the gradient
        tens = gwggrad(constC, hC1, hC2, T)

        T = sinkhorn(tens, a=p, b=q, eps=eps, max_iter=max_sinkhorn_iter)

        if cpt % 10 == 0:
            # we can speed up the process by checking for the error only all
            # the 10th iterations
            err = torch.linalg.norm(T - T_old)

        cpt += 1

    return T
//  [python]
def apply_distance_bound(data_region_i, data_region_i_orig, args):
    """
    Input:
        data_region_i: (S,3) tensor, current region i points
        data_region_i_orig: (S,3) tensor, original region i points
    Return:
        data_region_i: modified data_region_i
        count: number of points that exceed distance bound
    """
    with torch.no_grad():
        region_i_diff = data_region_i - data_region_i_orig #(S,3)
        region_i_diff_distance = torch.norm(region_i_diff, dim=1) #(S,)

        total_points = region_i_diff_distance.shape[0]
        count = 0
        for i in range(total_points):  # check bound
            if region_i_diff_distance[i] > args.dist_threshold:
                count += 1
                data_region_i[i].data = data_region_i_orig[i].data + args.dist_threshold * region_i_diff[i] / region_i_diff_distance[i]
    return data_region_i, count
//  [python]
def apply_distance_bound(data_region_i, data_region_i_orig, args):
    """
    Input:
        data_region_i: (S,3) tensor, current region i points
        data_region_i_orig: (S,3) tensor, original region i points
    Return:
        data_region_i: modified data_region_i
        count: number of points that exceed distance bound
    """
    with torch.no_grad():
        region_i_diff = data_region_i - data_region_i_orig #(S,3)
        region_i_diff_distance = torch.linalg.norm(region_i_diff, dim=1) #(S,)

        total_points = region_i_diff_distance.shape[0]
        count = 0
        for i in range(total_points):  # check bound
            if region_i_diff_distance[i] > args.dist_threshold:
                count += 1
                data_region_i[i].data = data_region_i_orig[i].data + args.dist_threshold * region_i_diff[i] / region_i_diff_distance[i]
    return data_region_i, count
//  [python]
def gradient_descent(data_region_i, objective, args):
    """
    Input:
        data_region_i: (S,3)
        objective: choice = ["inc", "dec"], indicate whether to perform gradient ascent or descent
    Return:
        (S,3), modified data_region_i
    """
    if_grad_none = False
    if data_region_i.grad != None:  # when all 3 lambdas exceed the bound, grad might be None. In this case, we do not update the region any more
        grad = data_region_i.grad.data
        norm = torch.norm(grad)
        delta = args.step * grad / norm if norm != 0 else 1e-8 # when norm=0, give a small perturbation
        data_region_i.data = data_region_i.data + delta if objective == "inc" else data_region_i.data - delta
    else:
        if_grad_none = True
    return data_region_i, if_grad_none
//  [python]
def gradient_descent(data_region_i, objective, args):
    """
    Input:
        data_region_i: (S,3)
        objective: choice = ["inc", "dec"], indicate whether to perform gradient ascent or descent
    Return:
        (S,3), modified data_region_i
    """
    if_grad_none = False
    if data_region_i.grad != None:  # when all 3 lambdas exceed the bound, grad might be None. In this case, we do not update the region any more
        grad = data_region_i.grad.data
        norm = torch.linalg.norm(grad)
        delta = args.step * grad / norm if norm != 0 else 1e-8 # when norm=0, give a small perturbation
        data_region_i.data = data_region_i.data + delta if objective == "inc" else data_region_i.data - delta
    else:
        if_grad_none = True
    return data_region_i, if_grad_none
//  [python]
def compute_error_accel(joints_gt, joints_pred, vis=None):
    """
    Computes acceleration error:
        1/(n-2) \sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}
    Note that for each frame that is not visible, three entries in the
    acceleration error should be zero'd out.
    Args:
        joints_gt (Nx14x3).
        joints_pred (Nx14x3).
        vis (N).
    Returns:
        error_accel (N-2).
    """
    # (N-2)x14x3
    accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
    accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]

    normed = torch.norm(accel_pred - accel_gt, dim=2)

    if vis is None:
        new_vis = torch.ones(len(normed), dtype=bool)
    else:
        invis = torch.logical_not(vis)
        invis1 = torch.roll(invis, -1)
        invis2 = torch.roll(invis, -2)
        new_invis = torch.logical_or(invis, np.logical_or(invis1, invis2))[:-2]
        new_vis = torch.logical_not(new_invis)

    acc=torch.mean(normed[new_vis], axis=1)

    return acc[~acc.isnan()]
//  [python]
def compute_error_accel(joints_gt, joints_pred, vis=None):
    """
    Computes acceleration error:
        1/(n-2) \sum_{i=1}^{n-1} X_{i-1} - 2X_i + X_{i+1}
    Note that for each frame that is not visible, three entries in the
    acceleration error should be zero'd out.
    Args:
        joints_gt (Nx14x3).
        joints_pred (Nx14x3).
        vis (N).
    Returns:
        error_accel (N-2).
    """
    # (N-2)x14x3
    accel_gt = joints_gt[:-2] - 2 * joints_gt[1:-1] + joints_gt[2:]
    accel_pred = joints_pred[:-2] - 2 * joints_pred[1:-1] + joints_pred[2:]

    normed = torch.linalg.norm(accel_pred - accel_gt, dim=2)

    if vis is None:
        new_vis = torch.ones(len(normed), dtype=bool)
    else:
        invis = torch.logical_not(vis)
        invis1 = torch.roll(invis, -1)
        invis2 = torch.roll(invis, -2)
        new_invis = torch.logical_or(invis, np.logical_or(invis1, invis2))[:-2]
        new_vis = torch.logical_not(new_invis)

    acc=torch.mean(normed[new_vis], axis=1)

    return acc[~acc.isnan()]
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4] 
            - src_keypts:    [bs, num_corr, 3]    
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:    
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 0.10: # for 3DMatch
            inlier_threshold_list = [0.10] * 20
        else: # for KITTI
            inlier_threshold_list = [1.2] * 20

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)
            L2_dis = torch.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1/(1 + (L2_dis/inlier_threshold)**2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4] 
            - src_keypts:    [bs, num_corr, 3]    
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:    
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        if self.inlier_threshold == 0.10: # for 3DMatch
            inlier_threshold_list = [0.10] * 20
        else: # for KITTI
            inlier_threshold_list = [1.2] * 20

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)
            L2_dis = torch.linalg.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1/(1 + (L2_dis/inlier_threshold)**2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def compute_angle(A, B):
    """
    Compute the angle between two tensors of the same size. The tensors will
     be flattened, after which the angle is computed.
    Args:
        A (torch.Tensor): First tensor
        B (torch.Tensor): Second tensor
    Returns: The angle between the two tensors in degrees
    """
    if contains_nan(A):
        print('tensor A contains nans:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans:')
        print(B)

    inner_product = torch.sum(A*B)  #equal to inner product of flattened tensors
    cosine = inner_product/(torch.norm(A, p='fro')*torch.norm(B, p='fro'))

    if cosine > 1 and cosine < 1 + 1e-5:
        cosine = torch.Tensor([1.])
    angle = 180/np.pi*torch.acos(cosine)
    if contains_nan(angle):
        print('angle computation causes NANs. cosines:')
        print(cosine)
    return angle
//  [python]
def compute_angle(A, B):
    """
    Compute the angle between two tensors of the same size. The tensors will
     be flattened, after which the angle is computed.
    Args:
        A (torch.Tensor): First tensor
        B (torch.Tensor): Second tensor
    Returns: The angle between the two tensors in degrees
    """
    if contains_nan(A):
        print('tensor A contains nans:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans:')
        print(B)

    inner_product = torch.sum(A*B)  #equal to inner product of flattened tensors
    cosine = inner_product/(torch.linalg.norm(A, ord='fro')*torch.linalg.norm(B, ord='fro'))

    if cosine > 1 and cosine < 1 + 1e-5:
        cosine = torch.Tensor([1.])
    angle = 180/np.pi*torch.acos(cosine)
    if contains_nan(angle):
        print('angle computation causes NANs. cosines:')
        print(cosine)
    return angle
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input: 
            - M:          [bs, num_corr, num_corr] the compatibility matrix 
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output: 
            - confidence  
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            # rank = torch.argsort(leading_eig, dim=1, descending=True)[:, 0:int(M.shape[1]*self.ratio)]
            # binary_sol = torch.zeros_like(leading_eig)
            # binary_sol[0, rank[0]] = 1
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input: 
            - M:          [bs, num_corr, num_corr] the compatibility matrix 
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output: 
            - confidence  
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.linalg.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            # rank = torch.argsort(leading_eig, dim=1, descending=True)[:, 0:int(M.shape[1]*self.ratio)]
            # binary_sol = torch.zeros_like(leading_eig)
            # binary_sol[0, rank[0]] = 1
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def step(self, closure=None):
        """
            Not supporting closure.
        """
        self._pre_step()
        self._partition_all_parameters()

        #checks for overflow, adjust the loss scale accordingly
        if self._overflow_check_and_loss_scale_update():
            if self.swap_optimizer:
                self.optimizer_swapper.log_timers()
            return

        norm_groups = self._get_norm_groups()
        scaled_global_grad_norm = torch.norm(torch.stack(norm_groups))

        # Stash unscaled gradient norm
        self._global_grad_norm = scaled_global_grad_norm / self.loss_scale

        timer_names = set()

        timer_names.add(OPTIMIZER_STEP_TIMER)
        self.timers(OPTIMIZER_STEP_TIMER).start()

        #update parameters one sub group at a time
        for sub_group_id, group in enumerate(self.fp16_groups):

            #prepare optimizer states, gradients and fp32 parameters for update
            self._prepare_sub_group(sub_group_id, timer_names)

            #scale the fp32 gradients
            self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm)

            #apply the optimizer step on the sub group and copy fp32 parameters to fp16
            self._optimizer_step(sub_group_id)

            #put fp16 parameters in appropriate location
            self._reassign_or_swap_out_partitioned_parameters(sub_group_id)

            #release memory or swap out optimizer states of fp32 parameters
            self._release_sub_group(sub_group_id, timer_names)

        self.timers(OPTIMIZER_STEP_TIMER).stop()

        self._post_step(timer_names)

        # warn user about caching allocator flushes
        memory_stats = get_accelerator().memory_stats()
        alloc_retries = memory_stats.get("num_alloc_retries")
        if alloc_retries is None:
            alloc_retries = 0
        if alloc_retries > self.n_caching_allocator_flushes:
            if dist.get_rank() == 0:
                logger.warning(
                    "%d pytorch allocator cache flushes since last step. this happens "
                    "when there is high memory pressure and is detrimental to "
                    "performance. if this is happening frequently consider adjusting "
                    "settings to reduce memory consumption. If you are unable to "
                    "make the cache flushes go away consider adding "
                    "get_accelerator().empty_cache() calls in your training loop to ensure "
                    "that all ranks flush their caches at the same time",
                    alloc_retries - self.n_caching_allocator_flushes)
            self.n_caching_allocator_flushes = alloc_retries
//  [python]
def step(self, closure=None):
        """
            Not supporting closure.
        """
        self._pre_step()
        self._partition_all_parameters()

        #checks for overflow, adjust the loss scale accordingly
        if self._overflow_check_and_loss_scale_update():
            if self.swap_optimizer:
                self.optimizer_swapper.log_timers()
            return

        norm_groups = self._get_norm_groups()
        scaled_global_grad_norm = torch.linalg.norm(torch.stack(norm_groups))

        # Stash unscaled gradient norm
        self._global_grad_norm = scaled_global_grad_norm / self.loss_scale

        timer_names = set()

        timer_names.add(OPTIMIZER_STEP_TIMER)
        self.timers(OPTIMIZER_STEP_TIMER).start()

        #update parameters one sub group at a time
        for sub_group_id, group in enumerate(self.fp16_groups):

            #prepare optimizer states, gradients and fp32 parameters for update
            self._prepare_sub_group(sub_group_id, timer_names)

            #scale the fp32 gradients
            self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm)

            #apply the optimizer step on the sub group and copy fp32 parameters to fp16
            self._optimizer_step(sub_group_id)

            #put fp16 parameters in appropriate location
            self._reassign_or_swap_out_partitioned_parameters(sub_group_id)

            #release memory or swap out optimizer states of fp32 parameters
            self._release_sub_group(sub_group_id, timer_names)

        self.timers(OPTIMIZER_STEP_TIMER).stop()

        self._post_step(timer_names)

        # warn user about caching allocator flushes
        memory_stats = get_accelerator().memory_stats()
        alloc_retries = memory_stats.get("num_alloc_retries")
        if alloc_retries is None:
            alloc_retries = 0
        if alloc_retries > self.n_caching_allocator_flushes:
            if dist.get_rank() == 0:
                logger.warning(
                    "%d pytorch allocator cache flushes since last step. this happens "
                    "when there is high memory pressure and is detrimental to "
                    "performance. if this is happening frequently consider adjusting "
                    "settings to reduce memory consumption. If you are unable to "
                    "make the cache flushes go away consider adding "
                    "get_accelerator().empty_cache() calls in your training loop to ensure "
                    "that all ranks flush their caches at the same time",
                    alloc_retries - self.n_caching_allocator_flushes)
            self.n_caching_allocator_flushes = alloc_retries
//  [python]
def calculate_jhmdb_PCK(predicted, gt, bbox, imgshape, thresh):

    # 0: neck    1:belly   2: face
    # 3: right shoulder  4: left shoulder
    # 5: right hip       6: left hip
    # 7: right elbow     8: left elbow
    # 9: right knee      10: left knee
    # 11: right wrist    12: left wrist
    # 13: right ankle    14: left ankle

    orderJHMDB = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    N = predicted.shape[0]

    HitPoint = torch.zeros((len(orderJHMDB))).to(gt.device)
    Point_to_use = torch.ones((len(orderJHMDB))).to(gt.device)

    
    test_gt = gt.reshape(-1, len(orderJHMDB), 2) * imgshape[0,[1,0]]
    test_out = predicted.reshape(-1, len(orderJHMDB), 2) * imgshape[0,[1,0]]

    seqError = torch.zeros(N, len(orderJHMDB)).to(gt.device)
    seqThresh = torch.zeros(N, len(orderJHMDB)).to(gt.device)

    bodysize = torch.max(bbox[:, 2] - bbox[:, 0], bbox[:, 3] - bbox[:, 1])

    for frame in range(0, N):
        gt_frame = test_gt[frame, :]
        pred_frame = test_out[frame, :]

        error_dis = torch.norm(gt_frame - pred_frame,
                                   p=2,
                                   dim=1,
                                   keepdim=False)

        seqError[frame] = error_dis
        seqThresh[frame] = (bodysize[frame] * thresh) * torch.ones(
                len(orderJHMDB)).to(gt.device)

    less_than_thresh = seqError <= seqThresh
    HitPoint = torch.sum(less_than_thresh, dim=0)

    finalPCK = torch.divide(HitPoint, N)

    return finalPCK
//  [python]
def calculate_jhmdb_PCK(predicted, gt, bbox, imgshape, thresh):

    # 0: neck    1:belly   2: face
    # 3: right shoulder  4: left shoulder
    # 5: right hip       6: left hip
    # 7: right elbow     8: left elbow
    # 9: right knee      10: left knee
    # 11: right wrist    12: left wrist
    # 13: right ankle    14: left ankle

    orderJHMDB = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
    N = predicted.shape[0]

    HitPoint = torch.zeros((len(orderJHMDB))).to(gt.device)
    Point_to_use = torch.ones((len(orderJHMDB))).to(gt.device)

    
    test_gt = gt.reshape(-1, len(orderJHMDB), 2) * imgshape[0,[1,0]]
    test_out = predicted.reshape(-1, len(orderJHMDB), 2) * imgshape[0,[1,0]]

    seqError = torch.zeros(N, len(orderJHMDB)).to(gt.device)
    seqThresh = torch.zeros(N, len(orderJHMDB)).to(gt.device)

    bodysize = torch.max(bbox[:, 2] - bbox[:, 0], bbox[:, 3] - bbox[:, 1])

    for frame in range(0, N):
        gt_frame = test_gt[frame, :]
        pred_frame = test_out[frame, :]

        error_dis = torch.linalg.norm(gt_frame - pred_frame,
                                   ord=2,
                                   dim=1,
                                   keepdim=False)

        seqError[frame] = error_dis
        seqThresh[frame] = (bodysize[frame] * thresh) * torch.ones(
                len(orderJHMDB)).to(gt.device)

    less_than_thresh = seqError <= seqThresh
    HitPoint = torch.sum(less_than_thresh, dim=0)

    finalPCK = torch.divide(HitPoint, N)

    return finalPCK
//  [python]
def test_rectangular_to_polar_dict(self, basic_stim):
        spc = po.simul.SteerablePyramidFreq(basic_stim.shape[-2:], height=5,
                                              order=1, is_complex=True, tight_frame=True).to(DEVICE)
        y = spc(basic_stim)
        energy, state = po.simul.non_linearities.rectangular_to_polar_dict(y, residuals=True)
        y_hat = po.simul.non_linearities.polar_to_rectangular_dict(energy, state, residuals=True)
        for key in y.keys():
            assert torch.norm(y[key] - y_hat[key]) < 1e-5
//  [python]
def test_rectangular_to_polar_dict(self, basic_stim):
        spc = po.simul.SteerablePyramidFreq(basic_stim.shape[-2:], height=5,
                                              order=1, is_complex=True, tight_frame=True).to(DEVICE)
        y = spc(basic_stim)
        energy, state = po.simul.non_linearities.rectangular_to_polar_dict(y, residuals=True)
        y_hat = po.simul.non_linearities.polar_to_rectangular_dict(energy, state, residuals=True)
        for key in y.keys():
            assert torch.linalg.norm(y[key] - y_hat[key]) < 1e-5
//  [python]
def dist(v1, v2, axis=None):
    """ Utility function to compute the Euclidean distance between two vectors.
    If only 1D vectors, an scalar is returned. If a 2D or 3D matrix is feed,
    the first dimension is interpreted as time and vector/matrix distance
    is computed along it."""

    if axis is None:
        if len(v1.shape) == 1: axis = 0        # normal 1D vectors
        elif len(v1.shape) == 2: axis = 1      # time x vector matrix, calculate along time
        elif len(v1.shape) == 3: axis = (1, 2)  # time x batch x vector matrix, calculate along time

    if isinstance(v1, torch.Tensor) and isinstance(v2, torch.Tensor):
        d = torch.norm(v1-v2, dim=axis, p=2).detach()
    elif isinstance(v1, np.ndarray) and isinstance(v2, np.ndarray):
        d = np.linalg.norm(v1-v2, axis=axis)
    else:
        raise ValueError('Invalid types {},{} for dist function'.format(
            type(v1), type(v2)))

    return d