jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, graph, feat, get_attention=False):
        # Check in degree and generate error
        if (graph.in_degrees() == 0).any():
            raise DGLError(
                "There are 0-in-degree nodes in the graph, "
                "output for those nodes will be invalid. "
                "This is harmful for some applications, "
                "causing silent performance regression. "
                "Adding self-loop on the input graph by "
                "calling `g = dgl.add_self_loop(g)` will resolve "
                "the issue. Setting ``allow_zero_in_degree`` "
                "to be `True` when constructing this module will "
                "suppress the check and let the code run."
            )
        # projection process to get importance vector y
        graph.ndata["y"] = torch.abs(
            torch.matmul(self.p, feat.T).view(-1)
        ) / torch.norm(self.p, p=2)
        # Use edge message passing function to get the weight from src node
        graph.apply_edges(fn.copy_u("y", "y"))
        # Select Top k neighbors
        subgraph = select_topk(graph.cpu(), self.k, "y").to(graph.device)
        # Sigmoid as information threshold
        subgraph.ndata["y"] = torch.sigmoid(subgraph.ndata["y"])
        # Using vector matrix elementwise mul for acceleration
        feat = subgraph.ndata["y"].view(-1, 1) * feat
        feat = self.feat_drop(feat)
        h = self.fc(feat).view(-1, self.num_heads, self.out_feats)
        el = (h * self.attn_l).sum(dim=-1).unsqueeze(-1)
        er = (h * self.attn_r).sum(dim=-1).unsqueeze(-1)
        # Assign the value on the subgraph
        subgraph.srcdata.update({"ft": h, "el": el})
        subgraph.dstdata.update({"er": er})
        # compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
        subgraph.apply_edges(fn.u_add_v("el", "er", "e"))
        e = self.leaky_relu(subgraph.edata.pop("e"))
        # compute softmax
        subgraph.edata["a"] = self.attn_drop(edge_softmax(subgraph, e))
        # message passing
        subgraph.update_all(fn.u_mul_e("ft", "a", "m"), fn.sum("m", "ft"))
        rst = subgraph.dstdata["ft"]
        # activation
        if self.activation:
            rst = self.activation(rst)
        # Residual
        if self.residual:
            rst = rst + self.residual_module(feat).view(
                feat.shape[0], -1, self.out_feats
            )

        if get_attention:
            return rst, subgraph.edata["a"]
        else:
            return rst
//  [python]
def forward(self, graph, feat, get_attention=False):
        # Check in degree and generate error
        if (graph.in_degrees() == 0).any():
            raise DGLError(
                "There are 0-in-degree nodes in the graph, "
                "output for those nodes will be invalid. "
                "This is harmful for some applications, "
                "causing silent performance regression. "
                "Adding self-loop on the input graph by "
                "calling `g = dgl.add_self_loop(g)` will resolve "
                "the issue. Setting ``allow_zero_in_degree`` "
                "to be `True` when constructing this module will "
                "suppress the check and let the code run."
            )
        # projection process to get importance vector y
        graph.ndata["y"] = torch.abs(
            torch.matmul(self.p, feat.T).view(-1)
        ) / torch.linalg.norm(self.p, ord=2)
        # Use edge message passing function to get the weight from src node
        graph.apply_edges(fn.copy_u("y", "y"))
        # Select Top k neighbors
        subgraph = select_topk(graph.cpu(), self.k, "y").to(graph.device)
        # Sigmoid as information threshold
        subgraph.ndata["y"] = torch.sigmoid(subgraph.ndata["y"])
        # Using vector matrix elementwise mul for acceleration
        feat = subgraph.ndata["y"].view(-1, 1) * feat
        feat = self.feat_drop(feat)
        h = self.fc(feat).view(-1, self.num_heads, self.out_feats)
        el = (h * self.attn_l).sum(dim=-1).unsqueeze(-1)
        er = (h * self.attn_r).sum(dim=-1).unsqueeze(-1)
        # Assign the value on the subgraph
        subgraph.srcdata.update({"ft": h, "el": el})
        subgraph.dstdata.update({"er": er})
        # compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
        subgraph.apply_edges(fn.u_add_v("el", "er", "e"))
        e = self.leaky_relu(subgraph.edata.pop("e"))
        # compute softmax
        subgraph.edata["a"] = self.attn_drop(edge_softmax(subgraph, e))
        # message passing
        subgraph.update_all(fn.u_mul_e("ft", "a", "m"), fn.sum("m", "ft"))
        rst = subgraph.dstdata["ft"]
        # activation
        if self.activation:
            rst = self.activation(rst)
        # Residual
        if self.residual:
            rst = rst + self.residual_module(feat).view(
                feat.shape[0], -1, self.out_feats
            )

        if get_attention:
            return rst, subgraph.edata["a"]
        else:
            return rst
//  [python]
def test_SYMEIG_random():
    m= 50
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D,U= SYMEIG.apply(M)
    assert( torch.norm(M-U@torch.diag(D)@U.t()) < D[0]*(m**2)*1e-14 )

    # since we always assume matrix M to be symmetric, the finite difference
    # perturbations should be symmetric as well
    M.requires_grad_(True)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        return SYMEIG.apply(M)
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_random():
    m= 50
    M= torch.rand(m, m, dtype=torch.float64)
    M= 0.5*(M+M.t())

    D,U= SYMEIG.apply(M)
    assert( torch.linalg.norm(M-U@torch.diag(D)@U.t()) < D[0]*(m**2)*1e-14 )

    # since we always assume matrix M to be symmetric, the finite difference
    # perturbations should be symmetric as well
    M.requires_grad_(True)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        return SYMEIG.apply(M)
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_3x3degenerate():
    M= torch.zeros((3,3),dtype=torch.float64)
    M[0,1]=M[0,2]=M[1,2]=1.
    M= 0.5*(M+M.t())
    print(M)

    D,U= SYMEIG.apply(M)
    assert( torch.norm(M-U@torch.diag(D)@U.t()) < D[0]*(M.size()[0]**2)*1e-14 )

    M.requires_grad_(True)
    torch.set_printoptions(precision=9)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        print(M)
        D,U= SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def test_SYMEIG_3x3degenerate():
    M= torch.zeros((3,3),dtype=torch.float64)
    M[0,1]=M[0,2]=M[1,2]=1.
    M= 0.5*(M+M.t())
    print(M)

    D,U= SYMEIG.apply(M)
    assert( torch.linalg.norm(M-U@torch.diag(D)@U.t()) < D[0]*(M.size()[0]**2)*1e-14 )

    M.requires_grad_(True)
    torch.set_printoptions(precision=9)
    def force_sym_eig(M):
        M=0.5*(M+M.t())
        print(M)
        D,U= SYMEIG.apply(M)
        return U
    assert(torch.autograd.gradcheck(force_sym_eig, M, eps=1e-6, atol=1e-4))
//  [python]
def estimate_A_lag_pytorch(n, n_aug, Y_h, Yh, rho, h, d, m, max_it, reg, verbose,return_res=0):
    bs, _, _ = Y_h.shape
    H = torch.zeros((bs, n_aug + (m - h), n_aug + (m - h)), dtype=Yh.dtype, device=Yh.device)
    H[:, :n_aug, :n_aug] = reg * batch_eye(bs=bs, n=n_aug, dtype=Yh.dtype, device=Yh.device)
    H[:, n_aug:, :n_aug] = Y_h.transpose(1, 2)
    H[:, :n_aug, n_aug:] = Y_h
    H[:, n_aug:, n_aug:] = - rho * batch_eye(n=m - h, bs=bs, dtype=Yh.dtype, device=Yh.device)
    H_inv = torch.inverse(H)
    Ats = []
    Yts = []
    At = torch.zeros((bs, n_aug, n), dtype=Yh.dtype, device=Yh.device)
    Yt = torch.zeros((bs, m - h, n), dtype=Yh.dtype, device=Yh.device)
    Ats.append(At)
    Yts.append(Yt)
    for i in range(max_it):
        if verbose:
            res = torch.norm(Ats[-1].transpose(1, 2) @ Y_h - Yh, np.inf)
            print("res: {}".format(res))
        rhs = torch.zeros((bs, n_aug + m - h, n), dtype=Yh.dtype, device=Yh.device)
        rhs[:, :n_aug, :] = -Y_h @ Yts[-1] - reg * Ats[-1]
        rhs[:, n_aug:, :] = Yh.transpose(1, 2) - Y_h.transpose(1, 2) @ Ats[-1]
        dZ = H_inv @ rhs
        Ats.append(Ats[-1] + dZ[:, :h * n + d, :])
        Yts.append(Yts[-1] + dZ[:, h * n + d:, :])
    A = Ats[-1].transpose(1, 2)
    if return_res:
        return A, res
    else:
        return A
//  [python]
def estimate_A_lag_pytorch(n, n_aug, Y_h, Yh, rho, h, d, m, max_it, reg, verbose,return_res=0):
    bs, _, _ = Y_h.shape
    H = torch.zeros((bs, n_aug + (m - h), n_aug + (m - h)), dtype=Yh.dtype, device=Yh.device)
    H[:, :n_aug, :n_aug] = reg * batch_eye(bs=bs, n=n_aug, dtype=Yh.dtype, device=Yh.device)
    H[:, n_aug:, :n_aug] = Y_h.transpose(1, 2)
    H[:, :n_aug, n_aug:] = Y_h
    H[:, n_aug:, n_aug:] = - rho * batch_eye(n=m - h, bs=bs, dtype=Yh.dtype, device=Yh.device)
    H_inv = torch.inverse(H)
    Ats = []
    Yts = []
    At = torch.zeros((bs, n_aug, n), dtype=Yh.dtype, device=Yh.device)
    Yt = torch.zeros((bs, m - h, n), dtype=Yh.dtype, device=Yh.device)
    Ats.append(At)
    Yts.append(Yt)
    for i in range(max_it):
        if verbose:
            res = torch.linalg.norm(Ats[-1].transpose(1, 2) @ Y_h - Yh, np.inf)
            print("res: {}".format(res))
        rhs = torch.zeros((bs, n_aug + m - h, n), dtype=Yh.dtype, device=Yh.device)
        rhs[:, :n_aug, :] = -Y_h @ Yts[-1] - reg * Ats[-1]
        rhs[:, n_aug:, :] = Yh.transpose(1, 2) - Y_h.transpose(1, 2) @ Ats[-1]
        dZ = H_inv @ rhs
        Ats.append(Ats[-1] + dZ[:, :h * n + d, :])
        Yts.append(Yts[-1] + dZ[:, h * n + d:, :])
    A = Ats[-1].transpose(1, 2)
    if return_res:
        return A, res
    else:
        return A
//  [python]
def frobenius_norm(U, VT):

    m, r = U.shape
    r, n = VT.shape
    if m * n * r < r * r * (m + n):
        return torch.norm(torch.matmul(U, VT))
    return torch.sqrt(torch.trace(torch.chain_matmul(VT, VT.T, U.T, U)))
//  [python]
def frobenius_norm(U, VT):

    m, r = U.shape
    r, n = VT.shape
    if m * n * r < r * r * (m + n):
        return torch.linalg.norm(torch.matmul(U, VT))
    return torch.sqrt(torch.trace(torch.chain_matmul(VT, VT.T, U.T, U)))
//  [python]
def non_orthogonality(U, VT):

    return [torch.norm(matrix - torch.diag(torch.diag(matrix))) / torch.norm(matrix)
            for matrix in [torch.matmul(U.T, U), torch.matmul(VT, VT.T)]]
//  [python]
def non_orthogonality(U, VT):

    return [torch.linalg.norm(matrix - torch.diag(torch.diag(matrix))) / torch.linalg.norm(matrix)
            for matrix in [torch.matmul(U.T, U), torch.matmul(VT, VT.T)]]
//  [python]
def __call__(self, X):
        N = X.size(0)
        dist = torch.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, p=2)  # [N, N]
        n_neighbors = min(self.knn + 1, N)

        knn_index = dist.topk(n_neighbors, -1, largest=False).indices[:, 1:]  # [N, knn]

        W = torch.zeros(N, N, device=X.device)
        W.scatter_(dim=-1, index=knn_index, value=1.0)

        return W
//  [python]
def __call__(self, X):
        N = X.size(0)
        dist = torch.linalg.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, ord=2)  # [N, N]
        n_neighbors = min(self.knn + 1, N)

        knn_index = dist.topk(n_neighbors, -1, largest=False).indices[:, 1:]  # [N, knn]

        W = torch.zeros(N, N, device=X.device)
        W.scatter_(dim=-1, index=knn_index, value=1.0)

        return W
//  [python]
def get_ordered_poses_and_k(
    poses: Float[Tensor, "num_poses 3 4"],
    Ks: Float[Tensor, "num_poses 3 3"],
) -> Tuple[Float[Tensor, "num_poses 3 4"], Float[Tensor, "num_poses 3 3"]]:
    """
    Returns ordered poses and intrinsics by euclidian distance between poses.

    Args:
        poses: list of camera poses
        Ks: list of camera intrinsics

    Returns:
        tuple of ordered poses and intrinsics

    """

    poses_num = len(poses)

    ordered_poses = torch.unsqueeze(poses[0], 0)
    ordered_ks = torch.unsqueeze(Ks[0], 0)

    # remove the first pose from poses
    poses = poses[1:]
    Ks = Ks[1:]

    for _ in range(poses_num - 1):
        distances = torch.norm(ordered_poses[-1][:, 3] - poses[:, :, 3], dim=1)
        idx = torch.argmin(distances)
        ordered_poses = torch.cat((ordered_poses, torch.unsqueeze(poses[idx], 0)), dim=0)
        ordered_ks = torch.cat((ordered_ks, torch.unsqueeze(Ks[idx], 0)), dim=0)
        poses = torch.cat((poses[0:idx], poses[idx + 1 :]), dim=0)
        Ks = torch.cat((Ks[0:idx], Ks[idx + 1 :]), dim=0)

    return ordered_poses, ordered_ks
//  [python]
def get_ordered_poses_and_k(
    poses: Float[Tensor, "num_poses 3 4"],
    Ks: Float[Tensor, "num_poses 3 3"],
) -> Tuple[Float[Tensor, "num_poses 3 4"], Float[Tensor, "num_poses 3 3"]]:
    """
    Returns ordered poses and intrinsics by euclidian distance between poses.

    Args:
        poses: list of camera poses
        Ks: list of camera intrinsics

    Returns:
        tuple of ordered poses and intrinsics

    """

    poses_num = len(poses)

    ordered_poses = torch.unsqueeze(poses[0], 0)
    ordered_ks = torch.unsqueeze(Ks[0], 0)

    # remove the first pose from poses
    poses = poses[1:]
    Ks = Ks[1:]

    for _ in range(poses_num - 1):
        distances = torch.linalg.norm(ordered_poses[-1][:, 3] - poses[:, :, 3], dim=1)
        idx = torch.argmin(distances)
        ordered_poses = torch.cat((ordered_poses, torch.unsqueeze(poses[idx], 0)), dim=0)
        ordered_ks = torch.cat((ordered_ks, torch.unsqueeze(Ks[idx], 0)), dim=0)
        poses = torch.cat((poses[0:idx], poses[idx + 1 :]), dim=0)
        Ks = torch.cat((Ks[0:idx], Ks[idx + 1 :]), dim=0)

    return ordered_poses, ordered_ks
//  [python]
def test_forward(self):
        module = GRN(in_channels=32, eps=1e-3)
        input_ = torch.rand(1, 28, 28, 32)
        gx = torch.norm(input_, p=2, dim=(1, 2), keepdim=True)
        nx = gx / (gx.mean(dim=3, keepdim=True) + 1e-3)
        expected_out = module.gamma * input_ * nx + module.beta + input_

        torch.testing.assert_allclose(
            module(input_, data_format='channel_last'), expected_out)

        input_ = input_.permute([0, 3, 1, 2])
        expected_out = expected_out.permute([0, 3, 1, 2])
        torch.testing.assert_allclose(
            module(input_, data_format='channel_first'), expected_out)
//  [python]
def test_forward(self):
        module = GRN(in_channels=32, eps=1e-3)
        input_ = torch.rand(1, 28, 28, 32)
        gx = torch.linalg.norm(input_, ord=2, dim=(1, 2), keepdim=True)
        nx = gx / (gx.mean(dim=3, keepdim=True) + 1e-3)
        expected_out = module.gamma * input_ * nx + module.beta + input_

        torch.testing.assert_allclose(
            module(input_, data_format='channel_last'), expected_out)

        input_ = input_.permute([0, 3, 1, 2])
        expected_out = expected_out.permute([0, 3, 1, 2])
        torch.testing.assert_allclose(
            module(input_, data_format='channel_first'), expected_out)
//  [python]
def ComputePower(mat_g, p,
                 iter_count=100,
                 error_tolerance=1e-6,
                 ridge_epsilon=1e-6):
  """A method to compute G^{-1/p} using a coupled Newton iteration.

  See for example equation 3.2 on page 9 of:
  A Schur-Newton Method for the Matrix p-th Root and its Inverse
  by Chun-Hua Guo and Nicholas J. Higham
  SIAM Journal on Matrix Analysis and Applications,
  2006, Vol. 28, No. 3 : pp. 788-804
  https://pdfs.semanticscholar.org/0abe/7f77433cf5908bfe2b79aa91af881da83858.pdf

  Args:
    mat_g: A square positive semidefinite matrix
    p: a positive integer
    iter_count: Stop iterating after this many rounds.
    error_tolerance: Threshold for stopping iteration
    ridge_epsilon: We add this times I to G, to make is positive definite.
                   For scaling, we multiply it by the largest eigenvalue of G.
  Returns:
    (mat_g + rI)^{-1/p} (r = ridge_epsilon * max_eigenvalue of mat_g).
  """
  shape = list(mat_g.shape)
  if len(shape) == 1:
    return torch.pow(mat_g + ridge_epsilon, -1/p)
  identity = torch.eye(shape[0], device=mat_g.get_device())
  if shape[0] == 1:
    return identity
  alpha = -1.0/p
  max_ev, _, _ = PowerIter(mat_g)
  ridge_epsilon *= max_ev
  mat_g += ridge_epsilon * identity
  z = (1 + p) / (2 * torch.norm(mat_g))
  # The best value for z is
  # (1 + p) * (c_max^{1/p} - c_min^{1/p}) /
  #            (c_max^{1+1/p} - c_min^{1+1/p})
  # where c_max and c_min are the largest and smallest singular values of
  # mat_g.
  # The above estimate assumes that c_max > c_min * 2^p
  # Can replace above line by the one below, but it is less accurate,
  # hence needs more iterations to converge.
  # z = (1 + p) / tf.trace(mat_g)
  # If we want the method to always converge, use z = 1 / norm(mat_g)
  # or z = 1 / tf.trace(mat_g), but these can result in many
  # extra iterations.

  mat_root = identity * torch.pow(z, 1.0/p)
  mat_m = mat_g * z
  error = torch.max(torch.abs(mat_m - identity))
  count = 0
  while error > error_tolerance and count < iter_count:
    tmp_mat_m = (1 - alpha) * identity + alpha * mat_m
    new_mat_root = torch.matmul(mat_root, tmp_mat_m)
    mat_m = torch.matmul(MatPower(tmp_mat_m, p), mat_m)
    new_error = torch.max(torch.abs(mat_m - identity))
    if new_error > error * 1.2:
      break
    mat_root = new_mat_root
    error = new_error
    count += 1
  return mat_root
//  [python]
def ComputePower(mat_g, p,
                 iter_count=100,
                 error_tolerance=1e-6,
                 ridge_epsilon=1e-6):
  """A method to compute G^{-1/p} using a coupled Newton iteration.

  See for example equation 3.2 on page 9 of:
  A Schur-Newton Method for the Matrix p-th Root and its Inverse
  by Chun-Hua Guo and Nicholas J. Higham
  SIAM Journal on Matrix Analysis and Applications,
  2006, Vol. 28, No. 3 : pp. 788-804
  https://pdfs.semanticscholar.org/0abe/7f77433cf5908bfe2b79aa91af881da83858.pdf

  Args:
    mat_g: A square positive semidefinite matrix
    p: a positive integer
    iter_count: Stop iterating after this many rounds.
    error_tolerance: Threshold for stopping iteration
    ridge_epsilon: We add this times I to G, to make is positive definite.
                   For scaling, we multiply it by the largest eigenvalue of G.
  Returns:
    (mat_g + rI)^{-1/p} (r = ridge_epsilon * max_eigenvalue of mat_g).
  """
  shape = list(mat_g.shape)
  if len(shape) == 1:
    return torch.pow(mat_g + ridge_epsilon, -1/p)
  identity = torch.eye(shape[0], device=mat_g.get_device())
  if shape[0] == 1:
    return identity
  alpha = -1.0/p
  max_ev, _, _ = PowerIter(mat_g)
  ridge_epsilon *= max_ev
  mat_g += ridge_epsilon * identity
  z = (1 + p) / (2 * torch.linalg.norm(mat_g))
  # The best value for z is
  # (1 + p) * (c_max^{1/p} - c_min^{1/p}) /
  #            (c_max^{1+1/p} - c_min^{1+1/p})
  # where c_max and c_min are the largest and smallest singular values of
  # mat_g.
  # The above estimate assumes that c_max > c_min * 2^p
  # Can replace above line by the one below, but it is less accurate,
  # hence needs more iterations to converge.
  # z = (1 + p) / tf.trace(mat_g)
  # If we want the method to always converge, use z = 1 / norm(mat_g)
  # or z = 1 / tf.trace(mat_g), but these can result in many
  # extra iterations.

  mat_root = identity * torch.pow(z, 1.0/p)
  mat_m = mat_g * z
  error = torch.max(torch.abs(mat_m - identity))
  count = 0
  while error > error_tolerance and count < iter_count:
    tmp_mat_m = (1 - alpha) * identity + alpha * mat_m
    new_mat_root = torch.matmul(mat_root, tmp_mat_m)
    mat_m = torch.matmul(MatPower(tmp_mat_m, p), mat_m)
    new_error = torch.max(torch.abs(mat_m - identity))
    if new_error > error * 1.2:
      break
    mat_root = new_mat_root
    error = new_error
    count += 1
  return mat_root
//  [python]
def get_cell_translations(self, **kwargs) -> Tuple[Tensor, Tensor, Tensor]:
        """Implement cell translation."""

        dd = {'dtype': self._dtype, 'device': self._device}

        # Extentions for lattice vectors
        pos_ext = kwargs.get('positive_extention', 1)
        neg_ext = kwargs.get('negative_extention', 1)

        # Ranges of cell translation on three dimentions
        _tmp = torch.floor(self.cutoff * torch.norm(
            self._invlatvec, dim=-1).T).T
        ranges = torch.stack([-(neg_ext + _tmp), pos_ext + _tmp])

        # For 1D/ 2D cell translation, non-periodic direction will be zero
        ranges[torch.stack([self._mask_zero, self._mask_zero])] = 0

        # length of the first, second and third column in ranges
        leng = ranges[1, :].long() - ranges[0, :].long() + 1

        # number of cells
        ncell = leng[..., 0] * leng[..., 1] * leng[..., 2]

        # Cell translation vectors in relative coordinates
        if not self._n_batch:  # -> single
            cellvec = torch.stack([
                torch.linspace(ranges[0, 0], ranges[1, 0],
                               leng[0], **dd).repeat_interleave(leng[2] * leng[1]),
                torch.linspace(ranges[0, 1], ranges[1, 1],
                               leng[1], **dd).repeat(leng[0]).repeat_interleave(leng[2]),
                torch.linspace(ranges[0, 2], ranges[1, 2],
                               leng[2], **dd).repeat(leng[0] * leng[1])])

        else:  # -> batch
            # Large values are padded at the end of short cell vectors.
            cellvec = pack([torch.stack([
                torch.linspace(iran[0, 0], iran[1, 0],
                               ile[0], **dd).repeat_interleave(ile[2] * ile[1]),
                torch.linspace(iran[0, 1], iran[1, 1],
                               ile[1], **dd).repeat(ile[0]).repeat_interleave(ile[2]),
                torch.linspace(iran[0, 2], iran[1, 2],
                               ile[2], **dd).repeat(ile[0] * ile[1])])
                            for ile, iran in zip(
                                    leng, ranges.transpose(-2, -3))], value=1e3)

        # Cell translation vectors in absolute units
        rcellvec = (torch.matmul(self.latvec.transpose(-1, -2),
                                 cellvec)).transpose(-1, -2)

        return cellvec.transpose(-1, -2), rcellvec, ncell
//  [python]
def get_cell_translations(self, **kwargs) -> Tuple[Tensor, Tensor, Tensor]:
        """Implement cell translation."""

        dd = {'dtype': self._dtype, 'device': self._device}

        # Extentions for lattice vectors
        pos_ext = kwargs.get('positive_extention', 1)
        neg_ext = kwargs.get('negative_extention', 1)

        # Ranges of cell translation on three dimentions
        _tmp = torch.floor(self.cutoff * torch.linalg.norm(
            self._invlatvec, dim=-1).T).T
        ranges = torch.stack([-(neg_ext + _tmp), pos_ext + _tmp])

        # For 1D/ 2D cell translation, non-periodic direction will be zero
        ranges[torch.stack([self._mask_zero, self._mask_zero])] = 0

        # length of the first, second and third column in ranges
        leng = ranges[1, :].long() - ranges[0, :].long() + 1

        # number of cells
        ncell = leng[..., 0] * leng[..., 1] * leng[..., 2]

        # Cell translation vectors in relative coordinates
        if not self._n_batch:  # -> single
            cellvec = torch.stack([
                torch.linspace(ranges[0, 0], ranges[1, 0],
                               leng[0], **dd).repeat_interleave(leng[2] * leng[1]),
                torch.linspace(ranges[0, 1], ranges[1, 1],
                               leng[1], **dd).repeat(leng[0]).repeat_interleave(leng[2]),
                torch.linspace(ranges[0, 2], ranges[1, 2],
                               leng[2], **dd).repeat(leng[0] * leng[1])])

        else:  # -> batch
            # Large values are padded at the end of short cell vectors.
            cellvec = pack([torch.stack([
                torch.linspace(iran[0, 0], iran[1, 0],
                               ile[0], **dd).repeat_interleave(ile[2] * ile[1]),
                torch.linspace(iran[0, 1], iran[1, 1],
                               ile[1], **dd).repeat(ile[0]).repeat_interleave(ile[2]),
                torch.linspace(iran[0, 2], iran[1, 2],
                               ile[2], **dd).repeat(ile[0] * ile[1])])
                            for ile, iran in zip(
                                    leng, ranges.transpose(-2, -3))], value=1e3)

        # Cell translation vectors in absolute units
        rcellvec = (torch.matmul(self.latvec.transpose(-1, -2),
                                 cellvec)).transpose(-1, -2)

        return cellvec.transpose(-1, -2), rcellvec, ncell
//  [python]
def axisangle2quat(vec, eps=1e-6):
    # type: (Tensor, float) -> Tensor
    # store input shape and reshape
    input_shape = vec.shape[:-1]
    vec = vec.reshape(-1, 3)

    # Grab angle
    angle = torch.norm(vec, dim=-1, keepdim=True)

    # Create return array
    quat = torch.zeros(torch.prod(torch.tensor(input_shape)), 4, device=vec.device)
    quat[:, 3] = 1.0

    # Grab indexes where angle is not zero an convert the input to its quaternion form
    idx = angle.reshape(-1) > eps
    quat[idx, :] = torch.cat([
        vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :],
        torch.cos(angle[idx, :] / 2.0)
    ], dim=-1)

    # Reshape and return output
    quat = quat.reshape(list(input_shape) + [4, ])
    return quat
//  [python]
def axisangle2quat(vec, eps=1e-6):
    # type: (Tensor, float) -> Tensor
    # store input shape and reshape
    input_shape = vec.shape[:-1]
    vec = vec.reshape(-1, 3)

    # Grab angle
    angle = torch.linalg.norm(vec, dim=-1, keepdim=True)

    # Create return array
    quat = torch.zeros(torch.prod(torch.tensor(input_shape)), 4, device=vec.device)
    quat[:, 3] = 1.0

    # Grab indexes where angle is not zero an convert the input to its quaternion form
    idx = angle.reshape(-1) > eps
    quat[idx, :] = torch.cat([
        vec[idx, :] * torch.sin(angle[idx, :] / 2.0) / angle[idx, :],
        torch.cos(angle[idx, :] / 2.0)
    ], dim=-1)

    # Reshape and return output
    quat = quat.reshape(list(input_shape) + [4, ])
    return quat
//  [python]
def cal_add_cuda(
        self, pred_RT, gt_RT, p3ds
    ):
        pred_p3ds = torch.mm(p3ds, pred_RT[:, :3].transpose(1, 0)) + pred_RT[:, 3]
        gt_p3ds = torch.mm(p3ds, gt_RT[:, :3].transpose(1, 0)) + gt_RT[:, 3]
        dis = torch.norm(pred_p3ds - gt_p3ds, dim=1)
        return torch.mean(dis)
//  [python]
def cal_add_cuda(
        self, pred_RT, gt_RT, p3ds
    ):
        pred_p3ds = torch.mm(p3ds, pred_RT[:, :3].transpose(1, 0)) + pred_RT[:, 3]
        gt_p3ds = torch.mm(p3ds, gt_RT[:, :3].transpose(1, 0)) + gt_RT[:, 3]
        dis = torch.linalg.norm(pred_p3ds - gt_p3ds, dim=1)
        return torch.mean(dis)
//  [python]
def cal_adds_cuda(
        self, pred_RT, gt_RT, p3ds
    ):
        N, _ = p3ds.size()
        pd = torch.mm(p3ds, pred_RT[:, :3].transpose(1, 0)) + pred_RT[:, 3]
        pd = pd.view(1, N, 3).repeat(N, 1, 1)
        gt = torch.mm(p3ds, gt_RT[:, :3].transpose(1, 0)) + gt_RT[:, 3]
        gt = gt.view(N, 1, 3).repeat(1, N, 1)
        dis = torch.norm(pd - gt, dim=2)
        mdis = torch.min(dis, dim=1)[0]
        return torch.mean(mdis)
//  [python]
def cal_adds_cuda(
        self, pred_RT, gt_RT, p3ds
    ):
        N, _ = p3ds.size()
        pd = torch.mm(p3ds, pred_RT[:, :3].transpose(1, 0)) + pred_RT[:, 3]
        pd = pd.view(1, N, 3).repeat(N, 1, 1)
        gt = torch.mm(p3ds, gt_RT[:, :3].transpose(1, 0)) + gt_RT[:, 3]
        gt = gt.view(N, 1, 3).repeat(1, N, 1)
        dis = torch.linalg.norm(pd - gt, dim=2)
        mdis = torch.min(dis, dim=1)[0]
        return torch.mean(mdis)
//  [python]
def frobenius_norm(self):
        A = torch.mm(
            self.data.view(-1, self.data.size(2)),
            self.data.view(-1, self.data.size(2)).t(),
        )
        return torch.norm(A)
//  [python]
def frobenius_norm(self):
        A = torch.mm(
            self.data.view(-1, self.data.size(2)),
            self.data.view(-1, self.data.size(2)).t(),
        )
        return torch.linalg.norm(A)
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        # check if the sequence is augmented by flipping
        self.is_mirror_extended = False
        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        # check if the sequence is augmented by flipping
        self.is_mirror_extended = False
        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.linalg.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def l2_norm(input, axis=1):
    norm = torch.norm(input, p=2, dim=axis, keepdim=True)
    output = torch.div(input, norm)
    return output
//  [python]
def l2_norm(input, axis=1):
    norm = torch.linalg.norm(input, ord=2, dim=axis, keepdim=True)
    output = torch.div(input, norm)
    return output
//  [python]
def fit(self, x, y, criterion=None):
        use_cuda = self.weight.data.is_cuda
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        reduction = criterion.reduction
        criterion.reduction = 'sum'
        if isinstance(x, np.ndarray) or isinstance(y, np.ndarray):
            x = torch.from_numpy(x)
            y = torch.from_numpy(y)
        if use_cuda:
            x = x.cuda()
            y = y.cuda()

        def eval_loss(w):
            w = w.reshape((self.out_features, -1))
            if self.weight.grad is not None:
                self.weight.grad = None
            if self.bias is None:
                self.weight.data.copy_(torch.from_numpy(w))
            else:
                if self.bias.grad is not None:
                    self.bias.grad = None
                self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
                self.bias.data.copy_(torch.from_numpy(w[:, -1]))
            y_pred = self(x).squeeze_(-1)
            loss = criterion(y_pred, y)
            loss.backward()
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    penalty = 0.5 * self.alpha * torch.norm(self.weight)**2
                elif self.penalty == "l1":
                    penalty = self.alpha * torch.norm(self.weight, p=1)
                    penalty.backward()
                loss = loss + penalty
            return loss.item()

        def eval_grad(w):
            dw = self.weight.grad.data
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    dw.add_(self.alpha, self.weight.data)
            if self.bias is not None:
                db = self.bias.grad.data
                dw = torch.cat((dw, db.view(-1, 1)), dim=1)
            return dw.cpu().numpy().ravel().astype("float64")

        w_init = self.weight.data
        if self.bias is not None:
            w_init = torch.cat((w_init, self.bias.data.view(-1, 1)), dim=1)
        w_init = w_init.cpu().numpy().astype("float64")

        w = optimize.fmin_l_bfgs_b(
            eval_loss, w_init, fprime=eval_grad, maxiter=100, disp=0)
        if isinstance(w, tuple):
            w = w[0]

        w = w.reshape((self.out_features, -1))
        self.weight.grad.data.zero_()
        if self.bias is None:
            self.weight.data.copy_(torch.from_numpy(w))
        else:
            self.bias.grad.data.zero_()
            self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
            self.bias.data.copy_(torch.from_numpy(w[:, -1]))
        criterion.reduction = reduction
//  [python]
def fit(self, x, y, criterion=None):
        use_cuda = self.weight.data.is_cuda
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        reduction = criterion.reduction
        criterion.reduction = 'sum'
        if isinstance(x, np.ndarray) or isinstance(y, np.ndarray):
            x = torch.from_numpy(x)
            y = torch.from_numpy(y)
        if use_cuda:
            x = x.cuda()
            y = y.cuda()

        def eval_loss(w):
            w = w.reshape((self.out_features, -1))
            if self.weight.grad is not None:
                self.weight.grad = None
            if self.bias is None:
                self.weight.data.copy_(torch.from_numpy(w))
            else:
                if self.bias.grad is not None:
                    self.bias.grad = None
                self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
                self.bias.data.copy_(torch.from_numpy(w[:, -1]))
            y_pred = self(x).squeeze_(-1)
            loss = criterion(y_pred, y)
            loss.backward()
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    penalty = 0.5 * self.alpha * torch.linalg.norm(self.weight)**2
                elif self.penalty == "l1":
                    penalty = self.alpha * torch.linalg.norm(self.weight, p=1)
                    penalty.backward()
                loss = loss + penalty
            return loss.item()

        def eval_grad(w):
            dw = self.weight.grad.data
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    dw.add_(self.alpha, self.weight.data)
            if self.bias is not None:
                db = self.bias.grad.data
                dw = torch.cat((dw, db.view(-1, 1)), dim=1)
            return dw.cpu().numpy().ravel().astype("float64")

        w_init = self.weight.data
        if self.bias is not None:
            w_init = torch.cat((w_init, self.bias.data.view(-1, 1)), dim=1)
        w_init = w_init.cpu().numpy().astype("float64")

        w = optimize.fmin_l_bfgs_b(
            eval_loss, w_init, fprime=eval_grad, maxiter=100, disp=0)
        if isinstance(w, tuple):
            w = w[0]

        w = w.reshape((self.out_features, -1))
        self.weight.grad.data.zero_()
        if self.bias is None:
            self.weight.data.copy_(torch.from_numpy(w))
        else:
            self.bias.grad.data.zero_()
            self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
            self.bias.data.copy_(torch.from_numpy(w[:, -1]))
        criterion.reduction = reduction
//  [python]
def test_basic_frobenius_norm(self):
        class SimpleOp(nn.Module):
            def forward(self, x):
                out = torch.norm(x, p="fro", dim=(0, 1), keepdim=True)
                return out
        x = torch.randn(2, 3, 4).float()
        self.checkExportImport(SimpleOp(), (x, ))
//  [python]
def test_basic_frobenius_norm(self):
        class SimpleOp(nn.Module):
            def forward(self, x):
                out = torch.linalg.norm(x, ord="fro", dim=(0, 1), keepdim=True)
                return out
        x = torch.randn(2, 3, 4).float()
        self.checkExportImport(SimpleOp(), (x, ))
//  [python]
def forward(self, x, gt):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).
            gt (Tensor): Ground-truth tensor with shape (n, c, h, w).

        Returns:
            Tensor: Forward results.
        """
        # extract vgg features
        x_features = self.vgg(x)
        gt_features = self.vgg(gt.detach())

        # calculate perceptual loss
        if self.perceptual_weight > 0:
            percep_loss = 0
            for k in x_features.keys():
                if self.criterion_type == 'fro':
                    percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
                else:
                    percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
            percep_loss *= self.perceptual_weight
        else:
            percep_loss = None

        # calculate style loss
        if self.style_weight > 0:
            style_loss = 0
            for k in x_features.keys():
                if self.criterion_type == 'fro':
                    style_loss += torch.norm(
                        self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
                else:
                    style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
                        gt_features[k])) * self.layer_weights[k]
            style_loss *= self.style_weight
        else:
            style_loss = None

        return percep_loss, style_loss
//  [python]
def forward(self, x, gt):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).
            gt (Tensor): Ground-truth tensor with shape (n, c, h, w).

        Returns:
            Tensor: Forward results.
        """
        # extract vgg features
        x_features = self.vgg(x)
        gt_features = self.vgg(gt.detach())

        # calculate perceptual loss
        if self.perceptual_weight > 0:
            percep_loss = 0
            for k in x_features.keys():
                if self.criterion_type == 'fro':
                    percep_loss += torch.linalg.norm(x_features[k] - gt_features[k], ord='fro') * self.layer_weights[k]
                else:
                    percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
            percep_loss *= self.perceptual_weight
        else:
            percep_loss = None

        # calculate style loss
        if self.style_weight > 0:
            style_loss = 0
            for k in x_features.keys():
                if self.criterion_type == 'fro':
                    style_loss += torch.linalg.norm(
                        self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), ord='fro') * self.layer_weights[k]
                else:
                    style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
                        gt_features[k])) * self.layer_weights[k]
            style_loss *= self.style_weight
        else:
            style_loss = None

        return percep_loss, style_loss
//  [python]
def linear_reg_loss(target: torch.Tensor,
                    feature: torch.Tensor,
                    reg: float):
    weight = fit_linear(target, feature, reg)
    pred = linear_reg_pred(feature, weight)
    return torch.norm((target - pred)) ** 2 + reg * torch.norm(weight) ** 2
//  [python]
def linear_reg_loss(target: torch.Tensor,
                    feature: torch.Tensor,
                    reg: float):
    weight = fit_linear(target, feature, reg)
    pred = linear_reg_pred(feature, weight)
    return torch.linalg.norm((target - pred)) ** 2 + reg * torch.linalg.norm(weight) ** 2
//  [python]
def linear_log_loss(target: torch.Tensor,
                    feature: torch.Tensor,
                    reg: float):
    weight = fit_linear(target, feature, reg)
    pred = linear_reg_pred(feature, weight)
    
    labels = logit(target)
    logits = logit(pred)
    return (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + reg * torch.norm(weight) ** 2
//  [python]
def linear_log_loss(target: torch.Tensor,
                    feature: torch.Tensor,
                    reg: float):
    weight = fit_linear(target, feature, reg)
    pred = linear_reg_pred(feature, weight)
    
    labels = logit(target)
    logits = logit(pred)
    return (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + reg * torch.linalg.norm(weight) ** 2
//  [python]
def compute_franka_reward(
        reset_buf, progress_buf, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf),
                            reset_buf)

    return rewards, reset_buf
//  [python]
def compute_franka_reward(
        reset_buf, progress_buf, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.linalg.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.linalg.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.linalg.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.linalg.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.linalg.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf),
                            reset_buf)

    return rewards, reset_buf
//  [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=torch.bool)
    else:
        invis = ~vis
        invis1 = torch.roll(invis, -1)
        invis2 = torch.roll(invis, -2)
        new_invis = invis | invis1 | invis2
        new_invis = new_invis[:-2]
        new_vis = ~new_invis
    return torch.mean(normed[new_vis], dim=1)
//  [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=torch.bool)
    else:
        invis = ~vis
        invis1 = torch.roll(invis, -1)
        invis2 = torch.roll(invis, -2)
        new_invis = invis | invis1 | invis2
        new_invis = new_invis[:-2]
        new_vis = ~new_invis
    return torch.mean(normed[new_vis], dim=1)
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
        assert pos.dim() == 2 and pos.size(1) == 2

        cart = pos[col] - pos[row]

        rho = torch.norm(cart, p=2, dim=-1).view(-1, 1)

        theta = torch.atan2(cart[..., 1], cart[..., 0]).view(-1, 1)
        theta = theta + (theta < 0).type_as(theta) * (2 * PI)

        if self.norm:
            rho = rho / (rho.max() if self.max is None else self.max)
            theta = theta / (2 * PI)

        polar = torch.cat([rho, theta], dim=-1)

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, polar.type_as(pos)], dim=-1)
        else:
            data.edge_attr = polar

        return data
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
        assert pos.dim() == 2 and pos.size(1) == 2

        cart = pos[col] - pos[row]

        rho = torch.linalg.norm(cart, p=2, dim=-1).view(-1, 1)

        theta = torch.atan2(cart[..., 1], cart[..., 0]).view(-1, 1)
        theta = theta + (theta < 0).type_as(theta) * (2 * PI)

        if self.norm:
            rho = rho / (rho.max() if self.max is None else self.max)
            theta = theta / (2 * PI)

        polar = torch.cat([rho, theta], dim=-1)

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, polar.type_as(pos)], dim=-1)
        else:
            data.edge_attr = polar

        return data