jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, activations, poses):

        self.W_reg = None
        if(self.share_W_ij):
            self.W_reg = torch.matmul(self.W_ij.squeeze(), self.W_ij.squeeze().permute(
                0, 1, 3, 2)) - torch.eye(4, device=self.device)
            self.W_reg = torch.norm(self.W_reg)

        if self.padding != 0:
            activations = F.pad(activations, self.padding)  # [1,1,1,1]
            poses = F.pad(poses, self.padding + [0]*4)  # [0,0,1,1,1,1]

        # share the matrices over (F*F), if class caps layer
        if self.share_W_ij:
            self.K = poses.shape[-1]  # out_caps (C) feature map size

        self.F = (poses.shape[-1] - self.K) // self.S + 1  # featuremap size

        # Out ← [?, B, P, P, F', F', K, K] ← [?, B, P, P, F, F]
        poses = poses.unfold(4, size=self.K, step=self.S).unfold(
            5, size=self.K, step=self.S)

        # Out ← [?, B, 1, P, P, 1, F', F', K, K] ← [?, B, P, P, F', F', K, K]
        poses = poses.unsqueeze(2).unsqueeze(5)

        # Out ← [?, B, F', F', K, K] ← [?, B, F, F]
        activations = activations.unfold(
            2, size=self.K, step=self.S).unfold(3, size=self.K, step=self.S)

        # Out ← [?, B, 1, 1, 1, F', F', K, K] ← [?, B, F', F', K, K]
        activations = activations.reshape(-1, self.B, 1,
                                          1, 1, *activations.shape[2:4], self.K, self.K)

        # Out ← [?, B, C, P, P, F', F', K, K] ← ([?, B, 1, P, P, 1, F', F', K, K] * [1, B, C, 1, P, P, 1, 1, K, K])
        V_ji = (poses * self.W_ij).sum(dim=4)  # matmul equiv.

        # Out ← [?, B, C, P*P, 1, F', F', K, K] ← [?, B, C, P, P, F', F', K, K]
        V_ji = V_ji.reshape(-1, self.B, self.C, self.P*self.P,
                            1, *V_ji.shape[-4:-2], self.K, self.K)

        if self.coor_add:
            # if class caps layer (featuremap size = 1)
            if V_ji.shape[-1] == 1:
                self.F = self.K  # 1->4

            # coordinates = torch.arange(self.F, dtype=torch.float32) / self.F
            coordinates = torch.arange(
                self.F, dtype=torch.float32).add(1.) / (self.F*10)
            i_vals = torch.zeros(self.P*self.P, self.F, 1, device=self.device)
            j_vals = torch.zeros(self.P*self.P, 1, self.F, device=self.device)
            i_vals[self.P-1, :, 0] = coordinates
            j_vals[2*self.P-1, 0, :] = coordinates

            if V_ji.shape[-1] == 1:  # if class caps layer
                # Out ← [?, B, C, P*P, 1, 1, 1, K=F, K=F] (class caps)
                V_ji = V_ji + (i_vals + j_vals).reshape(1, 1, 1,
                                                        self.P*self.P, 1, 1, 1, self.F, self.F)
                return activations, V_ji

            # Out ← [?, B, C, P*P, 1, F, F, K, K]
            V_ji = V_ji + (i_vals + j_vals).reshape(1, 1, 1,
                                                    self.P*self.P, 1, self.F, self.F, 1, 1)

        return activations, V_ji, self.W_reg
//  [python]
def forward(self, activations, poses):

        self.W_reg = None
        if(self.share_W_ij):
            self.W_reg = torch.matmul(self.W_ij.squeeze(), self.W_ij.squeeze().permute(
                0, 1, 3, 2)) - torch.eye(4, device=self.device)
            self.W_reg = torch.linalg.norm(self.W_reg)

        if self.padding != 0:
            activations = F.pad(activations, self.padding)  # [1,1,1,1]
            poses = F.pad(poses, self.padding + [0]*4)  # [0,0,1,1,1,1]

        # share the matrices over (F*F), if class caps layer
        if self.share_W_ij:
            self.K = poses.shape[-1]  # out_caps (C) feature map size

        self.F = (poses.shape[-1] - self.K) // self.S + 1  # featuremap size

        # Out ← [?, B, P, P, F', F', K, K] ← [?, B, P, P, F, F]
        poses = poses.unfold(4, size=self.K, step=self.S).unfold(
            5, size=self.K, step=self.S)

        # Out ← [?, B, 1, P, P, 1, F', F', K, K] ← [?, B, P, P, F', F', K, K]
        poses = poses.unsqueeze(2).unsqueeze(5)

        # Out ← [?, B, F', F', K, K] ← [?, B, F, F]
        activations = activations.unfold(
            2, size=self.K, step=self.S).unfold(3, size=self.K, step=self.S)

        # Out ← [?, B, 1, 1, 1, F', F', K, K] ← [?, B, F', F', K, K]
        activations = activations.reshape(-1, self.B, 1,
                                          1, 1, *activations.shape[2:4], self.K, self.K)

        # Out ← [?, B, C, P, P, F', F', K, K] ← ([?, B, 1, P, P, 1, F', F', K, K] * [1, B, C, 1, P, P, 1, 1, K, K])
        V_ji = (poses * self.W_ij).sum(dim=4)  # matmul equiv.

        # Out ← [?, B, C, P*P, 1, F', F', K, K] ← [?, B, C, P, P, F', F', K, K]
        V_ji = V_ji.reshape(-1, self.B, self.C, self.P*self.P,
                            1, *V_ji.shape[-4:-2], self.K, self.K)

        if self.coor_add:
            # if class caps layer (featuremap size = 1)
            if V_ji.shape[-1] == 1:
                self.F = self.K  # 1->4

            # coordinates = torch.arange(self.F, dtype=torch.float32) / self.F
            coordinates = torch.arange(
                self.F, dtype=torch.float32).add(1.) / (self.F*10)
            i_vals = torch.zeros(self.P*self.P, self.F, 1, device=self.device)
            j_vals = torch.zeros(self.P*self.P, 1, self.F, device=self.device)
            i_vals[self.P-1, :, 0] = coordinates
            j_vals[2*self.P-1, 0, :] = coordinates

            if V_ji.shape[-1] == 1:  # if class caps layer
                # Out ← [?, B, C, P*P, 1, 1, 1, K=F, K=F] (class caps)
                V_ji = V_ji + (i_vals + j_vals).reshape(1, 1, 1,
                                                        self.P*self.P, 1, 1, 1, self.F, self.F)
                return activations, V_ji

            # Out ← [?, B, C, P*P, 1, F, F, K, K]
            V_ji = V_ji + (i_vals + j_vals).reshape(1, 1, 1,
                                                    self.P*self.P, 1, self.F, self.F, 1, 1)

        return activations, V_ji, self.W_reg
//  [python]
def normalize_quaternion(q: torch.Tensor) -> torch.Tensor:
    return q / torch.norm(q, dim=-1, keepdim=True)
//  [python]
def normalize_quaternion(q: torch.Tensor) -> torch.Tensor:
    return q / torch.linalg.norm(q, dim=-1, keepdim=True)
//  [python]
def random_quaternions(n: int, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None):
    if isinstance(device, str):
        device = torch.device(device)
    q = torch.randn(n,4, device=device, dtype=dtype)

    return standardize_quaternion(q / torch.norm(q, dim=-1, keepdim=True))
//  [python]
def random_quaternions(n: int, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None):
    if isinstance(device, str):
        device = torch.device(device)
    q = torch.randn(n,4, device=device, dtype=dtype)

    return standardize_quaternion(q / torch.linalg.norm(q, dim=-1, keepdim=True))
//  [python]
def gsvd_torch(a, m, w):
    """
    :param a: Matrix to GSVD
    :param m: 1st Constraint, (u.T * m * u) = I
    :param w: 2nd Constraint, (v.T * w * v) = I
    :return: (u ,s, v)
    a = u * diag(s) * v.T  st. (u.T * m * u) = I, and (v.T * w * v) = I
    """

    (aHeight, aWidth) = a.shape
    (mHeight, mWidth) = m.shape
    (wHeight, wWidth) = w.shape

    assert(aHeight == mHeight), "Height not consistent."
    assert(aWidth == wWidth), "Weight not consistent"
    mSqrt = sqrtm_torch(m)
    wSqrt = sqrtm_torch(w)

    mSqrtInv = mSqrt.inverse()
    wSqrtInv = wSqrt.inverse()

    _a = torch.mm(torch.mm(mSqrt, a), wSqrt)

    (_u, _s, _v) = _a.svd()

    u = torch.mm(mSqrtInv, _u)
    v = torch.mm(wSqrtInv, _v)
    s = _s
    print('\tGSVD: distance {:2.6f}, matrix a norm {:2.6f}.'.
          format(torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())).cpu().detach().numpy(), tol * torch.norm(a).cpu().detach().numpy()))
    assert torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())) <= (tol * torch.norm(a)), "The result is not the GSVD."
    return u, s, v.t()
//  [python]
def gsvd_torch(a, m, w):
    """
    :param a: Matrix to GSVD
    :param m: 1st Constraint, (u.T * m * u) = I
    :param w: 2nd Constraint, (v.T * w * v) = I
    :return: (u ,s, v)
    a = u * diag(s) * v.T  st. (u.T * m * u) = I, and (v.T * w * v) = I
    """

    (aHeight, aWidth) = a.shape
    (mHeight, mWidth) = m.shape
    (wHeight, wWidth) = w.shape

    assert(aHeight == mHeight), "Height not consistent."
    assert(aWidth == wWidth), "Weight not consistent"
    mSqrt = sqrtm_torch(m)
    wSqrt = sqrtm_torch(w)

    mSqrtInv = mSqrt.inverse()
    wSqrtInv = wSqrt.inverse()

    _a = torch.mm(torch.mm(mSqrt, a), wSqrt)

    (_u, _s, _v) = _a.svd()

    u = torch.mm(mSqrtInv, _u)
    v = torch.mm(wSqrtInv, _v)
    s = _s
    print('\tGSVD: distance {:2.6f}, matrix a norm {:2.6f}.'.
          format(torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())).cpu().detach().numpy(), tol * torch.linalg.norm(a).cpu().detach().numpy()))
    assert torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())) <= (tol * torch.linalg.norm(a)), "The result is not the GSVD."
    return u, s, v.t()
//  [python]
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    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 first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    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(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    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 first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    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(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def pred(xy):
    '''
    

    Parameters
    ----------
    输入xytensor，然而输出位移，都是GPU的张量
    '''
    out = torch.zeros((len(xy), 1))
    out[torch.norm(xy, dim=1)<r0] = model_h1(xy[torch.norm(xy, dim=1)<r0]) # 内部的位移场，不需要RBF距离函数以及特解，只需要一个神经网络即可 
    out[torch.norm(xy, dim=1)>=r0] = model_p(xy[torch.norm(xy, dim=1)>=r0]) + RBF(xy[torch.norm(xy, dim=1)>=r0]) * model_h2(xy[torch.norm(xy, dim=1)>=r0]) # 外部的神经网络
    return out
//  [python]
def pred(xy):
    '''
    

    Parameters
    ----------
    输入xytensor，然而输出位移，都是GPU的张量
    '''
    out = torch.zeros((len(xy), 1))
    out[torch.linalg.norm(xy, dim=1)<r0] = model_h1(xy[torch.linalg.norm(xy, dim=1)<r0]) # 内部的位移场，不需要RBF距离函数以及特解，只需要一个神经网络即可 
    out[torch.linalg.norm(xy, dim=1)>=r0] = model_p(xy[torch.linalg.norm(xy, dim=1)>=r0]) + RBF(xy[torch.linalg.norm(xy, dim=1)>=r0]) * model_h2(xy[torch.linalg.norm(xy, dim=1)>=r0]) # 外部的神经网络
    return out
//  [python]
def gsvd_torch_acc(a, m, w):
    """
    :param a: Matrix to GSVD
    :param m: 1st Constraint, (u.T * m * u) = I
    :param w: 2nd Constraint, (v.T * w * v) = I
    :return: (u ,s, v)
    a = u * diag(s) * v.T  st. (u.T * m * u) = I, and (v.T * w * v) = I
    """

    (aHeight, aWidth) = a.shape
    (mHeight, mWidth) = m.shape
    (wHeight, wWidth) = w.shape

    assert(aHeight == mHeight), "Height not consistent."
    assert(aWidth == wWidth), "Weight not consistent"
    mSqrt = sqrtm_torch_acc(m)
    wSqrt = sqrtm_torch_acc(w)

    mSqrtInv = mSqrt.inverse()
    wSqrtInv = wSqrt.inverse()

    _a = torch.mm(torch.mm(mSqrt, a), wSqrt)

    (_u, _s, _v) = _a.svd()

    u = torch.mm(mSqrtInv, _u)
    v = torch.mm(wSqrtInv, _v)
    s = _s
    print('\tGSVD: distance {:2.6f}, matrix a norm {:2.6f}.'.
          format(torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())).cpu().detach().numpy(), tol * torch.norm(a).cpu().detach().numpy()))
    assert torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())) <= (tol * torch.norm(a)), "The result is not the GSVD."
    return u, s, v.t()
//  [python]
def gsvd_torch_acc(a, m, w):
    """
    :param a: Matrix to GSVD
    :param m: 1st Constraint, (u.T * m * u) = I
    :param w: 2nd Constraint, (v.T * w * v) = I
    :return: (u ,s, v)
    a = u * diag(s) * v.T  st. (u.T * m * u) = I, and (v.T * w * v) = I
    """

    (aHeight, aWidth) = a.shape
    (mHeight, mWidth) = m.shape
    (wHeight, wWidth) = w.shape

    assert(aHeight == mHeight), "Height not consistent."
    assert(aWidth == wWidth), "Weight not consistent"
    mSqrt = sqrtm_torch_acc(m)
    wSqrt = sqrtm_torch_acc(w)

    mSqrtInv = mSqrt.inverse()
    wSqrtInv = wSqrt.inverse()

    _a = torch.mm(torch.mm(mSqrt, a), wSqrt)

    (_u, _s, _v) = _a.svd()

    u = torch.mm(mSqrtInv, _u)
    v = torch.mm(wSqrtInv, _v)
    s = _s
    print('\tGSVD: distance {:2.6f}, matrix a norm {:2.6f}.'.
          format(torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())).cpu().detach().numpy(), tol * torch.linalg.norm(a).cpu().detach().numpy()))
    assert torch.dist(a, torch.mm(torch.mm(u, s.diag()), v.t())) <= (tol * torch.linalg.norm(a)), "The result is not the GSVD."
    return u, s, v.t()
//  [python]
def chamfer_distance(pc, recons_pc):
    '''
    Parameters
    ----------
    pc              Input point cloud
    recons_pc       Reconstructed point cloud

    Returns Shape loss -> how far the reconstructed points (PC) are estimated from the input point cloud
    -------

    '''
    # pdb.set_trace()

    pred_to_gt = torch.cat([torch.squeeze(
        torch.norm(pc[i].unsqueeze(1) - recons_pc[i].unsqueeze(0), dim=2, p=None).topk(1, largest=False, dim=0)[
            0]) for i in range(len(recons_pc))], dim=0)
    gt_to_pred = torch.cat([torch.squeeze(
        torch.norm(recons_pc[i].unsqueeze(1) - pc[i].unsqueeze(0), dim=2, p=None).topk(1, largest=False, dim=0)[
            0]) for i in range(len(pc))], dim=0)

    return torch.mean(pred_to_gt) + torch.mean(gt_to_pred)
//  [python]
def chamfer_distance(pc, recons_pc):
    '''
    Parameters
    ----------
    pc              Input point cloud
    recons_pc       Reconstructed point cloud

    Returns Shape loss -> how far the reconstructed points (PC) are estimated from the input point cloud
    -------

    '''
    # pdb.set_trace()

    pred_to_gt = torch.cat([torch.squeeze(
        torch.linalg.norm(pc[i].unsqueeze(1) - recons_pc[i].unsqueeze(0), dim=2, ord=None).topk(1, largest=False, dim=0)[
            0]) for i in range(len(recons_pc))], dim=0)
    gt_to_pred = torch.cat([torch.squeeze(
        torch.linalg.norm(recons_pc[i].unsqueeze(1) - pc[i].unsqueeze(0), dim=2, ord=None).topk(1, largest=False, dim=0)[
            0]) for i in range(len(pc))], dim=0)

    return torch.mean(pred_to_gt) + torch.mean(gt_to_pred)
//  [python]
def shape_loss(pc, kp):
    '''
    Parameters
    ----------
    pc      Input point cloud
    kp      Estimated key-points

    Returns Shape loss -> how far the key-points are estimated from the input point cloud
    -------

    '''
    loss = torch.cat([torch.squeeze(
        torch.norm(pc[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None).topk(1, largest=False, dim=0)[
            0]) for i in range(len(kp))], dim=0)
    return torch.mean(loss)
//  [python]
def shape_loss(pc, kp):
    '''
    Parameters
    ----------
    pc      Input point cloud
    kp      Estimated key-points

    Returns Shape loss -> how far the key-points are estimated from the input point cloud
    -------

    '''
    loss = torch.cat([torch.squeeze(
        torch.linalg.norm(pc[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2).topk(1, largest=False, dim=0)[
            0]) for i in range(len(kp))], dim=0)
    return torch.mean(loss)
//  [python]
def pre_logits(self, x):
        x = x / torch.norm(x, p=2, dim=1, keepdim=True)
        return x
//  [python]
def pre_logits(self, x):
        x = x / torch.linalg.norm(x, ord=2, dim=1, keepdim=True)
        return x
//  [python]
def separation_loss(kp):
    '''
    Parameters
    ----------
    kp:         Key-points
    Method:     compute distances of each point from all the points in "kp"
                consider minimum two distances (distance of a point form itself (distance==0) and the next closest (distance>0))
                take mean of the distances from the closest point (distance>0)

    Returns     separation loss ->  average distance of every point from closest points
    -------
    '''
    min_distances = torch.cat([torch.squeeze(
        torch.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None).topk(2, largest=False, dim=0)[
            0]) for i in range(len(kp))], dim=0)

    return 1/torch.mean(min_distances[min_distances>0])
//  [python]
def separation_loss(kp):
    min_distances = torch.cat([torch.squeeze(
        torch.linalg.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, ord=None).topk(2, largest=False, dim=0)[
            0]) for i in range(len(kp))], dim=0)

    return 1/torch.mean(min_distances[min_distances>0])
//  [python]
def normalize_over_channels(x):
    channel_norms = torch.norm(x, dim=1, keepdim=True)
    return x / channel_norms
//  [python]
def normalize_over_channels(x):
    channel_norms = torch.linalg.norm(x, dim=1, keepdim=True)
    return x / channel_norms
//  [python]
def loss(self, feat, target, **kwargs):
        losses = dict()
        # compute loss
        if self.with_len:
            loss = self.compute_loss(feat, target, m_norm2=torch.norm(target, p=2, dim=1))
        else:
            loss = self.compute_loss(feat, target)
        losses['loss'] = loss
        return losses
//  [python]
def loss(self, feat, target, **kwargs):
        losses = dict()
        # compute loss
        if self.with_len:
            loss = self.compute_loss(feat, target, m_norm2=torch.linalg.norm(target, ord=2, dim=1))
        else:
            loss = self.compute_loss(feat, target)
        losses['loss'] = loss
        return losses
//  [python]
def overlap_loss(kp, threshold=0.05):
    '''
    Parameters
    ----------
    kp:         Key-points
    threshold   allowable overlap between the key-points
    Method:     Find distance of every point from all the points
                select the minimum distances that are greater than 0 (distance from itself)
                return count of the separated distances => final loss

    Returns     separation loss -> avoid estimation of multiple key-points on the same 3D location
    -------
    '''

    distances = torch.cat([torch.squeeze(
        torch.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None)) for i in range(len(kp))], dim=0)

    return torch.sum(distances[(distances < threshold)] >0).float() / len(distances)*len(distances)
//  [python]
def overlap_loss(kp, threshold=0.05):
    '''
    Parameters
    ----------
    kp:         Key-points
    threshold   allowable overlap between the key-points
    Method:     Find distance of every point from all the points
                select the minimum distances that are greater than 0 (distance from itself)
                return count of the separated distances => final loss

    Returns     separation loss -> avoid estimation of multiple key-points on the same 3D location
    -------
    '''

    distances = torch.cat([torch.squeeze(
        torch.linalg.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None)) for i in range(len(kp))], dim=0)

    return torch.sum(distances[(distances < threshold)] >0).float() / len(distances)*len(distances)
//  [python]
def overlap_loss_torch_error(kp, threshold=0.05):
    '''
    Parameters
    ----------
    kp:         Key-points
    threshold   allowable overlap between the key-points
    Method:     Find distance of every point from all the points
                select the minimum distances that are greater than 0 (distance from itself)
                return count of the separated distances => final loss

    Returns     separation loss -> avoid estimation of multiple key-points on the same 3D location
    -------
    '''

    distances = torch.cat([torch.squeeze(
        torch.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None)) for i in range(len(kp))], dim=0)

    return torch.count_nonzero(distances[(distances < threshold)] >0) / len(kp)*len(kp)
//  [python]
def overlap_loss_torch_error(kp, threshold=0.05):
    '''
    Parameters
    ----------
    kp:         Key-points
    threshold   allowable overlap between the key-points
    Method:     Find distance of every point from all the points
                select the minimum distances that are greater than 0 (distance from itself)
                return count of the separated distances => final loss

    Returns     separation loss -> avoid estimation of multiple key-points on the same 3D location
    -------
    '''

    distances = torch.cat([torch.squeeze(
        torch.linalg.norm(kp[i].unsqueeze(1) - kp[i].unsqueeze(0), dim=2, p=None)) for i in range(len(kp))], dim=0)

    return torch.count_nonzero(distances[(distances < threshold)] >0) / len(kp)*len(kp)
//  [python]
def test_mds(n=100, d=10):
    X = torch.randn(n, d)
    X = X / torch.norm(X, dim=-1, keepdim=True) * 0.9
    X = from_poincare(X)
    # print(X.shape)
    D = distance(X.unsqueeze(-2), X.unsqueeze(-3))
    # print(D.shape)
    # print(D-D.transpose(0,1))

    X_ = mds(D, d)
    # print(X_.shape)
    D_ = distance(X_.unsqueeze(-2), X_.unsqueeze(-3))
    print(D - D_)
//  [python]
def test_mds(n=100, d=10):
    X = torch.randn(n, d)
    X = X / torch.linalg.norm(X, dim=-1, keepdim=True) * 0.9
    X = from_poincare(X)
    # print(X.shape)
    D = distance(X.unsqueeze(-2), X.unsqueeze(-3))
    # print(D.shape)
    # print(D-D.transpose(0,1))

    X_ = mds(D, d)
    # print(X_.shape)
    D_ = distance(X_.unsqueeze(-2), X_.unsqueeze(-3))
    print(D - D_)
//  [python]
def calculate_adaptive_weight(self, recon_loss, g_loss, last_layer, disc_weight_max):
        recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0]
        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]

        d_weight = torch.norm(recon_grads) / (torch.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, disc_weight_max).detach()
        return d_weight
//  [python]
def calculate_adaptive_weight(self, recon_loss, g_loss, last_layer, disc_weight_max):
        recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0]
        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]

        d_weight = torch.linalg.norm(recon_grads) / (torch.linalg.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, disc_weight_max).detach()
        return d_weight
//  [python]
def axis(self):
        r"""
        If a None axis is given, then the default axis is the
        normalised velocity.

        Returns:
            (N,d) Tensor

        """
        if self.axis_is_none:
            return self.vel / torch.norm(self.vel, dim=1).reshape((self.N, 1))
        else:
            return self.__axis
//  [python]
def axis(self):
        r"""
        If a None axis is given, then the default axis is the
        normalised velocity.

        Returns:
            (N,d) Tensor

        """
        if self.axis_is_none:
            return self.vel / torch.linalg.norm(self.vel, dim=1).reshape((self.N, 1))
        else:
            return self.__axis
//  [python]
def order_parameter(self):
        r"""Compute the norm of the average velocity.

        Returns:
            float:

            The order parameter.

        """
        vec_phi = (1. / self.N) * self.vel.sum(0)
        phi = torch.norm(vec_phi)
        return float(phi)
//  [python]
def order_parameter(self):
        r"""Compute the norm of the average velocity.

        Returns:
            float:

            The order parameter.

        """
        vec_phi = (1. / self.N) * self.vel.sum(0)
        phi = torch.linalg.norm(vec_phi)
        return float(phi)
//  [python]
def cos_similarity(emb_a: torch.Tensor, emb_b: torch.Tensor, eps=torch.tensor(3.5e-4)) -> torch.Tensor:
    """
    Calculate cosine similarities of the given two set of tensors. The output is an N by N
    matrix where N is the number of feature vectors.

    Args:
        a (Tensor):
            Matrix containing speaker representation vectors. (N x embedding_dim)
        b (Tensor):
            Matrix containing speaker representation vectors. (N x embedding_dim)

    Returns:
        res (Tensor):
            N by N matrix containing the cosine similarities of the values.
    """
    # If number of embedding count is 1, it creates nan values
    if emb_a.shape[0] == 1 or emb_b.shape[0] == 1:
        raise ValueError(f"Number of feature vectors should be greater than 1 but got {emb_a.shape} and {emb_b.shape}")
    a_norm = emb_a / (torch.norm(emb_a, dim=1).unsqueeze(1) + eps)
    b_norm = emb_b / (torch.norm(emb_b, dim=1).unsqueeze(1) + eps)
    res = torch.mm(a_norm, b_norm.transpose(0, 1))
    res.fill_diagonal_(1)
    return res
//  [python]
def cos_similarity(emb_a: torch.Tensor, emb_b: torch.Tensor, eps=torch.tensor(3.5e-4)) -> torch.Tensor:
    """
    Calculate cosine similarities of the given two set of tensors. The output is an N by N
    matrix where N is the number of feature vectors.

    Args:
        a (Tensor):
            Matrix containing speaker representation vectors. (N x embedding_dim)
        b (Tensor):
            Matrix containing speaker representation vectors. (N x embedding_dim)

    Returns:
        res (Tensor):
            N by N matrix containing the cosine similarities of the values.
    """
    # If number of embedding count is 1, it creates nan values
    if emb_a.shape[0] == 1 or emb_b.shape[0] == 1:
        raise ValueError(f"Number of feature vectors should be greater than 1 but got {emb_a.shape} and {emb_b.shape}")
    a_norm = emb_a / (torch.linalg.norm(emb_a, dim=1).unsqueeze(1) + eps)
    b_norm = emb_b / (torch.linalg.norm(emb_b, dim=1).unsqueeze(1) + eps)
    res = torch.mm(a_norm, b_norm.transpose(0, 1))
    res.fill_diagonal_(1)
    return res
//  [python]
def test_loss_LOO(Xs, ys, test_X, test_y):
    M = len(Xs)
    D = Xs[0].shape[1]

    Xs = torch.cat(Xs).reshape(-1, D)
    ys = torch.cat(ys).reshape(-1, 1)


    pinv = torch.pinverse(Xs)
    test_loss = (torch.norm( test_X @ pinv @ ys - test_y )) ** 2

    loo_values = []
    loo_losses = []
    for i in range(M):
        loo_dataset = []
        loo_label = []

        for j, (dataset, label) in enumerate(zip(Xs, ys)):
            if i == j: continue

            loo_dataset.append(dataset)
            loo_label.append(label)
        
        loo_dataset = torch.cat(loo_dataset).reshape(-1, D)
        loo_label = torch.cat(loo_label).reshape(-1, 1)

        pinv = torch.pinverse(loo_dataset)

        loo_test_loss = (torch.norm(  test_X @ pinv @ loo_label - test_y ))**2
        loo_losses.append(loo_test_loss.item())
        loo_values.append((loo_test_loss -  test_loss).item())

    # print('------Leave One Out Statistics ------')
    # print("Full test loss:", test_loss.item())
    # print("Leave-one-out test losses:", loo_losses)
    # print("Leave-one-out (loo_loss - full_loss) :", loo_values)
    # print('-------------------------------------')

    return np.asarray(loo_values)
//  [python]
def test_loss_LOO(Xs, ys, test_X, test_y):
    M = len(Xs)
    D = Xs[0].shape[1]

    Xs = torch.cat(Xs).reshape(-1, D)
    ys = torch.cat(ys).reshape(-1, 1)


    pinv = torch.pinverse(Xs)
    test_loss = (torch.linalg.norm( test_X @ pinv @ ys - test_y )) ** 2

    loo_values = []
    loo_losses = []
    for i in range(M):
        loo_dataset = []
        loo_label = []

        for j, (dataset, label) in enumerate(zip(Xs, ys)):
            if i == j: continue

            loo_dataset.append(dataset)
            loo_label.append(label)
        
        loo_dataset = torch.cat(loo_dataset).reshape(-1, D)
        loo_label = torch.cat(loo_label).reshape(-1, 1)

        pinv = torch.pinverse(loo_dataset)

        loo_test_loss = (torch.linalg.norm(  test_X @ pinv @ loo_label - test_y ))**2
        loo_losses.append(loo_test_loss.item())
        loo_values.append((loo_test_loss -  test_loss).item())

    # print('------Leave One Out Statistics ------')
    # print("Full test loss:", test_loss.item())
    # print("Leave-one-out test losses:", loo_losses)
    # print("Leave-one-out (loo_loss - full_loss) :", loo_values)
    # print('-------------------------------------')

    return np.asarray(loo_values)
//  [python]
def theta(self):
        r"""Polar angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.norm(Omega)
        return float(torch.acos(Omega[2]))
//  [python]
def theta(self):
        r"""Polar angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.linalg.norm(Omega)
        return float(torch.acos(Omega[2]))
//  [python]
def phi(self):
        r"""Azimuthal angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.norm(Omega)
        phi_pi = float(torch.atan2(Omega[1], Omega[0]))
        return np.mod(phi_pi, 2 * np.pi)
//  [python]
def phi(self):
        r"""Azimuthal angle of the direction of motion"""
        Omega = self.vel.sum(0)
        Omega = Omega / torch.linalg.norm(Omega)
        phi_pi = float(torch.atan2(Omega[1], Omega[0]))
        return np.mod(phi_pi, 2 * np.pi)
//  [python]
def normalised(self, kappa, who=None, with_who=None, **kwargs):
        r"""The normalised average velocity.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the normalised average velocity
        around particle :math:`(x_i,v_i)` is

        .. math::

            \Omega_i = \frac{\sum_j K(x_i-x_j) v_j}{\left|\sum_j K(x_i-x_j) v_j\right|},

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The normalised averaged velocity multiplied by the
            concentration parameter.

        """
        J, = self.linear_local_average(self.vel, who=who, with_who=with_who)
        Omega = J / torch.norm(J, dim=1).reshape((J.shape[0], 1))
        return kappa * Omega
//  [python]
def normalised(self, kappa, who=None, with_who=None, **kwargs):
        r"""The normalised average velocity.

        Given a system of particles with positions :math:`(x_i)_i` and
        velocities :math:`(v_i)_i`, the normalised average velocity
        around particle :math:`(x_i,v_i)` is

        .. math::

            \Omega_i = \frac{\sum_j K(x_i-x_j) v_j}{\left|\sum_j K(x_i-x_j) v_j\right|},

        where :math:`K` is the interaction kernel.

        Args:
            kappa ((1,) Tensor or (M,) Tensor): The concentration
                parameter. If the size of kappa is bigger than 1 then
                its size must be equal to the number of True values in
                **who**.
            who ((N,) BoolTensor or None, optional): The rows of
                the interaction kernel is ``pos[who,:]`` (or
                :attr:`pos` if **who** is None). Default is None.
            with_who ((N,) BoolTensor or None, optional): The
                columns of the interaction kernel is
                ``pos[with_who,:]`` (or :attr:`pos` if **with_who** is
                None). Default is None.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            Tensor:

            The normalised averaged velocity multiplied by the
            concentration parameter.

        """
        J, = self.linear_local_average(self.vel, who=who, with_who=with_who)
        Omega = J / torch.linalg.norm(J, dim=1).reshape((J.shape[0], 1))
        return kappa * Omega
//  [python]
def local_vel(self, xmin=0, xmax=1, ymin=0, ymax=1, zmin=0, zmax=1,
                  who=None):
        if who is None:
            index = (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax)
        else:
            index = who
        Omega = self.vel[index, :].sum(0)
        order_parameter = torch.norm(Omega) / len(self.vel[index, 0])
        Omega = Omega / torch.norm(Omega)
        theta = float(torch.acos(Omega[2]))
        phi = np.mod(float(torch.atan2(Omega[1], Omega[0])), 2 * np.pi)
        return order_parameter, theta, phi
//  [python]
def local_vel(self, xmin=0, xmax=1, ymin=0, ymax=1, zmin=0, zmax=1,
                  who=None):
        if who is None:
            index = (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax) \
                    and (self.pos[:, 0] >= xmin and self.pos[:, 0] <= xmax)
        else:
            index = who
        Omega = self.vel[index, :].sum(0)
        order_parameter = torch.linalg.norm(Omega) / len(self.vel[index, 0])
        Omega = Omega / torch.linalg.norm(Omega)
        theta = float(torch.acos(Omega[2]))
        phi = np.mod(float(torch.atan2(Omega[1], Omega[0])), 2 * np.pi)
        return order_parameter, theta, phi
//  [python]
def norm(self, X, G):
        return torch.norm(torch.solve(G, X)[0])
//  [python]
def norm(self, X, G):
        return torch.linalg.norm(torch.solve(G, X)[0])
//  [python]
def bounded_angular_velocity(self, target, angvel_max, dt,
                                 who=None, **kwargs):
        r"""Bounded angle between the velocity and the target.

        The maximum angle between the velocity and the new target cannot
        exceed the value ``angvel_max * dt``. The output tensor
        **new_target** is of the form::

            new_target = a * vel + b * target

        where :math:`a,b>0` are such that the cosine between
        **new_target** and :attr:`vel` is equal to::

            max_cos := max(cos(angvel_max * dt), cos_target)

        where **cos_target** is the cosine between **target** and
        :attr:`vel`. The solution is::

            b = (1 - max_cos ** 2) / (1 - cos_target ** 2)
            a = |target| / |vel| * (max_cos - b * cos_target)

        Args:
            target ((M,d) Tensor): Target.
            angvel_max (float): Maximum angular velocity.
            dt (float): Interval of time.
            who ((N,) BoolTensor, optional): The velocities of the
                particles associated to the target are ``vel[who,:]``
                (or :attr:`vel` if **who** is None). Default is None.
            **kwargs: Arbitrary keywords arguments.

        Returns:
            (M,d) Tensor:

            The new target.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        M = target.shape[0]
        vel = self.vel[who, :]
        norm_vel = torch.norm(vel, dim=1).reshape((M, 1))
        norm_target = torch.norm(target, dim=1).reshape((M, 1))

        cos_max = math.cos(angvel_max * dt)
        dot = (vel * target).sum(axis=1).reshape((M, 1))
        cos = dot / (norm_vel * norm_target)
        max_cos = dot / (norm_vel * norm_target)
        max_cos[max_cos <= cos_max] = cos_max

        # Add .01 for stability
        b = torch.sqrt((1.01 - max_cos ** 2) / (1.01 - cos ** 2))
        a = (norm_target / norm_vel) * (max_cos - b * cos)

        new_target = a * vel + b * target

        return new_target
//  [python]
def bounded_angular_velocity(self, target, angvel_max, dt,
                                 who=None, **kwargs):
        r"""Bounded angle between the velocity and the target.

        The maximum angle between the velocity and the new target cannot
        exceed the value ``angvel_max * dt``. The output tensor
        **new_target** is of the form::

            new_target = a * vel + b * target

        where :math:`a,b>0` are such that the cosine between
        **new_target** and :attr:`vel` is equal to::

            max_cos := max(cos(angvel_max * dt), cos_target)

        where **cos_target** is the cosine between **target** and
        :attr:`vel`. The solution is::

            b = (1 - max_cos ** 2) / (1 - cos_target ** 2)
            a = |target| / |vel| * (max_cos - b * cos_target)

        Args:
            target ((M,d) Tensor): Target.
            angvel_max (float): Maximum angular velocity.
            dt (float): Interval of time.
            who ((N,) BoolTensor, optional): The velocities of the
                particles associated to the target are ``vel[who,:]``
                (or :attr:`vel` if **who** is None). Default is None.
            **kwargs: Arbitrary keywords arguments.

        Returns:
            (M,d) Tensor:

            The new target.

        """
        if who is None:
            who = torch.ones(self.N, dtype=torch.bool, device=self.pos.device)
        M = target.shape[0]
        vel = self.vel[who, :]
        norm_vel = torch.linalg.norm(vel, dim=1).reshape((M, 1))
        norm_target = torch.linalg.norm(target, dim=1).reshape((M, 1))

        cos_max = math.cos(angvel_max * dt)
        dot = (vel * target).sum(axis=1).reshape((M, 1))
        cos = dot / (norm_vel * norm_target)
        max_cos = dot / (norm_vel * norm_target)
        max_cos[max_cos <= cos_max] = cos_max

        # Add .01 for stability
        b = torch.sqrt((1.01 - max_cos ** 2) / (1.01 - cos ** 2))
        a = (norm_target / norm_vel) * (max_cos - b * cos)

        new_target = a * vel + b * target

        return new_target
//  [python]
def forward(self,
        ipts2d: torch.Tensor, # N, 2
        ipts3d: torch.Tensor, # M, 2
        ) -> torch.Tensor:
        num_pts = len(ipts2d)
        p2d = ipts2d[None, :, :] # B, N, 2
        p3d = ipts3d[None, :, :] # B, N, 2
        # compute corss dist
        p2d_dist = torch.norm((p2d[:, :, None, :] - p2d[:, None, :, :]), dim=-1)
        p3d_dist = torch.norm((p3d[:, :, None, :] - p3d[:, None, :, :]), dim=-1)
        cross_dist = torch.abs(p2d_dist - p3d_dist)
        # compute first order measure
        SC_measure = torch.clamp(1.0 - cross_dist ** 2 / self.d_thre ** 2, min=0)
        hard_SC_measure = (cross_dist < self.d_thre).float()
        # select reliable seed correspondences
        confidence = self.cal_leading_eigenvector(SC_measure, method='power')
        import pdb
        pdb.set_trace()
        seeds = self.pick_seeds(p2d_dist, confidence, R=self.nms_radius, max_num=int(num_pts * self.ratio))
        # compute second order measure
        SC2_dist_thre = self.d_thre / 2
        hard_SC_measure_tight = (cross_dist < SC2_dist_thre).float()
        seed_hard_SC_measure = hard_SC_measure.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_pts))
        seed_hard_SC_measure_tight = hard_SC_measure_tight.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_pts))
        SC2_measure = torch.matmul(seed_hard_SC_measure_tight, hard_SC_measure_tight) * seed_hard_SC_measure
        return confidence
//  [python]
def forward(self,
        ipts2d: torch.Tensor, # N, 2
        ipts3d: torch.Tensor, # M, 2
        ) -> torch.Tensor:
        num_pts = len(ipts2d)
        p2d = ipts2d[None, :, :] # B, N, 2
        p3d = ipts3d[None, :, :] # B, N, 2
        # compute corss dist
        p2d_dist = torch.linalg.norm((p2d[:, :, None, :] - p2d[:, None, :, :]), dim=-1)
        p3d_dist = torch.linalg.norm((p3d[:, :, None, :] - p3d[:, None, :, :]), dim=-1)
        cross_dist = torch.abs(p2d_dist - p3d_dist)
        # compute first order measure
        SC_measure = torch.clamp(1.0 - cross_dist ** 2 / self.d_thre ** 2, min=0)
        hard_SC_measure = (cross_dist < self.d_thre).float()
        # select reliable seed correspondences
        confidence = self.cal_leading_eigenvector(SC_measure, method='power')
        import pdb
        pdb.set_trace()
        seeds = self.pick_seeds(p2d_dist, confidence, R=self.nms_radius, max_num=int(num_pts * self.ratio))
        # compute second order measure
        SC2_dist_thre = self.d_thre / 2
        hard_SC_measure_tight = (cross_dist < SC2_dist_thre).float()
        seed_hard_SC_measure = hard_SC_measure.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_pts))
        seed_hard_SC_measure_tight = hard_SC_measure_tight.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_pts))
        SC2_measure = torch.matmul(seed_hard_SC_measure_tight, hard_SC_measure_tight) * seed_hard_SC_measure
        return confidence
//  [python]
def update(self):
        """Set and update iteration variables."""
        if self.ivars["istep"] == 0:
            X_norm = float(torch.norm(self.X))
            iX_norm = X_norm**-1
            A_norm = float(torch.norm(_utils.matmul(self.A, self.X))) * iX_norm
            B_norm = float(torch.norm(_utils.matmul(self.B, self.X))) * iX_norm
            self.fvars["X_norm"] = X_norm
            self.fvars["A_norm"] = A_norm
            self.fvars["B_norm"] = B_norm
            self.ivars["iterations_left"] = self.iparams["niter"]
            self.ivars["converged_count"] = 0
            self.ivars["converged_end"] = 0

        if self.method == "ortho":
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars["iterations_left"] = self.ivars["iterations_left"] - 1
        self.ivars["istep"] = self.ivars["istep"] + 1
//  [python]
def update(self):
        """Set and update iteration variables."""
        if self.ivars["istep"] == 0:
            X_norm = float(torch.linalg.norm(self.X))
            iX_norm = X_norm**-1
            A_norm = float(torch.linalg.norm(_utils.matmul(self.A, self.X))) * iX_norm
            B_norm = float(torch.linalg.norm(_utils.matmul(self.B, self.X))) * iX_norm
            self.fvars["X_norm"] = X_norm
            self.fvars["A_norm"] = A_norm
            self.fvars["B_norm"] = B_norm
            self.ivars["iterations_left"] = self.iparams["niter"]
            self.ivars["converged_count"] = 0
            self.ivars["converged_end"] = 0

        if self.method == "ortho":
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars["iterations_left"] = self.ivars["iterations_left"] - 1
        self.ivars["istep"] = self.ivars["istep"] + 1
//  [python]
def norm_vector(self, vec, p=2, dim=0):
        """
        Implements the normalize() function from sklearn
        """
        vec_norm = torch.norm(vec, p=p, dim=dim)
        return vec.div(vec_norm.expand_as(vec))
//  [python]
def norm_vector(self, vec, p=2, dim=0):
        """
        Implements the normalize() function from sklearn
        """
        vec_norm = torch.linalg.norm(vec, ord=p, dim=dim)
        return vec.div(vec_norm.expand_as(vec))
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(x, (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)), mode='reflect')
        x = x.squeeze(1)
        x = torch.view_as_real(
            torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False, return_complex=True)
        )  # [B, F, TT, 2] (Note: torch.stft() returns complex tensor [B, F, TT]; converted via view_as_real)
        mag = torch.norm(x, p=2, dim=-1)  # [B, F, TT]

        return mag
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(x, (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)), mode='reflect')
        x = x.squeeze(1)
        x = torch.view_as_real(
            torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False, return_complex=True)
        )  # [B, F, TT, 2] (Note: torch.stft() returns complex tensor [B, F, TT]; converted via view_as_real)
        mag = torch.linalg.norm(x, ord=2, dim=-1)  # [B, F, TT]

        return mag
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].

        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars["converged_count"]
        tol = self.fparams["tol"]
        A_norm = self.fvars["A_norm"]
        B_norm = self.fvars["B_norm"]
        E, X, R = self.E, self.X, self.R
        rerr = (
            torch.norm(R, 2, (0,))
            * (torch.norm(X, 2, (0,)) * (A_norm + E[: X.shape[-1]] * B_norm)) ** -1
        )
        converged = rerr < tol
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert (
            count >= prev_count
        ), f"the number of converged eigenpairs (was {prev_count}, got {count}) cannot decrease"
        self.ivars["converged_count"] = count
        self.tvars["rerr"] = rerr
        return count
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].

        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars["converged_count"]
        tol = self.fparams["tol"]
        A_norm = self.fvars["A_norm"]
        B_norm = self.fvars["B_norm"]
        E, X, R = self.E, self.X, self.R
        rerr = (
            torch.linalg.norm(R, 2, (0,))
            * (torch.linalg.norm(X, 2, (0,)) * (A_norm + E[: X.shape[-1]] * B_norm)) ** -1
        )
        converged = rerr < tol
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert (
            count >= prev_count
        ), f"the number of converged eigenpairs (was {prev_count}, got {count}) cannot decrease"
        self.ivars["converged_count"] = count
        self.tvars["rerr"] = rerr
        return count