jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def cheirality_check(models, points):
    epipoles = torch.linalg.cross(models[..., 0, :], models[..., 2, :])
    degenerate_mask = torch.linalg.norm(epipoles, dim=-1) < 1e-6
    if torch.any(degenerate_mask):
        epi2 = torch.linalg.cross(models[degenerate_mask][..., 1, :], models[degenerate_mask][..., 2, :])
        epipoles[degenerate_mask] = epi2
    s1 = models[..., 0, 0].unsqueeze(1) * points[..., 2] + models[..., 1, 0].unsqueeze(1) * points[..., 3] + models[..., 2, 0].unsqueeze(1)
    s2 = epipoles[..., 1].unsqueeze(1) - epipoles[..., 2].unsqueeze(1) * points[..., 1]
    s = s1 * s2 > 0

    not_ok = s.float().mean(dim=-1) > 0.5
    s[not_ok] = ~s[not_ok]
    return s
//  [python]
def lossFunction(xvPairs):
    stretchPenalty = 0
    for (i,(x_i,v_i)) in enumerate(xvPairs):
        x_fixed1 = x_i[attachmentIdx[0]*3:attachmentIdx[0]*3+3]
        x_fixed2 = x_i[attachmentIdx[1]*3:attachmentIdx[1]*3+3]
        clip_dist = torch.norm(x_fixed2-x_fixed1)
        stretchPenalty += torch.clamp(torch.abs(clip_dist - CLIP_REST_DIST) - 1.0, min=0.0, max=None) * 0.2

    directionPenalty, targetLoss, targetLast = 0, 0, 0
    for (x_last,v_i) in  xvPairs:
        for (i, (idx1, idx2)) in enumerate(CLIP_DIR_VERTEX_PAIR):
            dir = x_last[idx1*3:idx1*3+3]-x_last[idx2*3:idx2*3+3]
            dirGoal = targetShape[idx1*3:idx1*3+3]- targetShape[idx2*3:idx2*3+3]
            dirGoaltorch =  common.toTorchTensor(dirGoal, False, False)
            cosine = torch.clamp(torch.dot(torch.nn.functional.normalize(dir, dim=0), torch.nn.functional.normalize(dirGoaltorch, dim=0)),  max=0.5, min=None)
            directionPenalty += (0.5 - cosine) * 3.0

        targetLoss += torch.nn.functional.smooth_l1_loss(x_last,  targetshape_torch)
    targetLast = torch.nn.functional.smooth_l1_loss(xvPairs[-1][0],  targetshape_torch)
    succeed = targetLast < 1.0

    loss = {'succeed': succeed, 'target' : targetLoss,'stretch' : stretchPenalty, 'direction' : directionPenalty, 'total' : stretchPenalty + targetLoss + directionPenalty}
    return loss
//  [python]
def lossFunction(xvPairs):
    stretchPenalty = 0
    for (i,(x_i,v_i)) in enumerate(xvPairs):
        x_fixed1 = x_i[attachmentIdx[0]*3:attachmentIdx[0]*3+3]
        x_fixed2 = x_i[attachmentIdx[1]*3:attachmentIdx[1]*3+3]
        clip_dist = torch.linalg.norm(x_fixed2-x_fixed1)
        stretchPenalty += torch.clamp(torch.abs(clip_dist - CLIP_REST_DIST) - 1.0, min=0.0, max=None) * 0.2

    directionPenalty, targetLoss, targetLast = 0, 0, 0
    for (x_last,v_i) in  xvPairs:
        for (i, (idx1, idx2)) in enumerate(CLIP_DIR_VERTEX_PAIR):
            dir = x_last[idx1*3:idx1*3+3]-x_last[idx2*3:idx2*3+3]
            dirGoal = targetShape[idx1*3:idx1*3+3]- targetShape[idx2*3:idx2*3+3]
            dirGoaltorch =  common.toTorchTensor(dirGoal, False, False)
            cosine = torch.clamp(torch.dot(torch.nn.functional.normalize(dir, dim=0), torch.nn.functional.normalize(dirGoaltorch, dim=0)),  max=0.5, min=None)
            directionPenalty += (0.5 - cosine) * 3.0

        targetLoss += torch.nn.functional.smooth_l1_loss(x_last,  targetshape_torch)
    targetLast = torch.nn.functional.smooth_l1_loss(xvPairs[-1][0],  targetshape_torch)
    succeed = targetLast < 1.0

    loss = {'succeed': succeed, 'target' : targetLoss,'stretch' : stretchPenalty, 'direction' : directionPenalty, 'total' : stretchPenalty + targetLoss + directionPenalty}
    return loss
//  [python]
def normalize(self, matches):
        dev = matches.device
        # The number of points in each minimal sample
        num_points = matches.shape[1]
        # Calculate the mass point for each minimal sample
        mass = torch.mean(matches, dim=1)
        # Substract the mass point of each minimal sample from the corresponding points in both images
        matches = matches - torch.unsqueeze(mass, 1).repeat(1, num_points, 1)
        # Calculate the distances from the mass point for each minimal sample in the source image
        distances1 = torch.norm(matches[:, :, :2], dim=2)
        # Calculate the distances from the mass point for each minimal sample in the destination image
        distances2 = torch.norm(matches[:, :, 2:], dim=2)
        # Calculate the average distances in the source image
        avg_distance1 = torch.mean(distances1, dim=1)
        # Calculate the average distances in the destination image
        avg_distance2 = torch.mean(distances2, dim=1)
        # Calculate the scaling to make the average distances sqrt(2) in the source image
        ratio1 = math.sqrt(2) / avg_distance1
        # Calculate the scaling to make the average distances sqrt(2) in the destination image
        ratio2 = math.sqrt(2) / avg_distance2

        # Calculate the normalized matches in the source image
        normalized_matches1 = matches[:, :, :2] * ratio1.view(-1, 1, 1).repeat(1, num_points, 2)
        # Calculate the normalized matches in the destination image
        normalized_matches2 = matches[:, :, 2:] * ratio2.view(-1, 1, 1).repeat(1, num_points, 2)

        # Initialize the normalizing transformations for each minimal sample in the source image
        T1 = torch.zeros((matches.shape[0], 3, 3), device=dev)
        # Initialize the normalizing transformations for each minimal sample in the destination image
        T2 = torch.zeros((matches.shape[0], 3, 3), device=dev)

        # Calculate the transformation parameters
        T1[:, 0, 0] = T1[:, 1, 1] = ratio1[:]
        T2[:, 0, 0] = T2[:, 1, 1] = ratio2[:]
        T1[:, 2, 2] = T2[:, 2, 2] = 1
        T1[:, 0, 2] = -ratio1 * mass[:, 0]
        T1[:, 1, 2] = -ratio1 * mass[:, 1]
        T2[:, 2, 0] = -ratio2 * mass[:, 2]
        T2[:, 2, 1] = -ratio2 * mass[:, 3]

        return torch.cat((normalized_matches1, normalized_matches2), dim=2), T1, T2
//  [python]
def normalize(self, matches):
        dev = matches.device
        # The number of points in each minimal sample
        num_points = matches.shape[1]
        # Calculate the mass point for each minimal sample
        mass = torch.mean(matches, dim=1)
        # Substract the mass point of each minimal sample from the corresponding points in both images
        matches = matches - torch.unsqueeze(mass, 1).repeat(1, num_points, 1)
        # Calculate the distances from the mass point for each minimal sample in the source image
        distances1 = torch.linalg.norm(matches[:, :, :2], dim=2)
        # Calculate the distances from the mass point for each minimal sample in the destination image
        distances2 = torch.linalg.norm(matches[:, :, 2:], dim=2)
        # Calculate the average distances in the source image
        avg_distance1 = torch.mean(distances1, dim=1)
        # Calculate the average distances in the destination image
        avg_distance2 = torch.mean(distances2, dim=1)
        # Calculate the scaling to make the average distances sqrt(2) in the source image
        ratio1 = math.sqrt(2) / avg_distance1
        # Calculate the scaling to make the average distances sqrt(2) in the destination image
        ratio2 = math.sqrt(2) / avg_distance2

        # Calculate the normalized matches in the source image
        normalized_matches1 = matches[:, :, :2] * ratio1.view(-1, 1, 1).repeat(1, num_points, 2)
        # Calculate the normalized matches in the destination image
        normalized_matches2 = matches[:, :, 2:] * ratio2.view(-1, 1, 1).repeat(1, num_points, 2)

        # Initialize the normalizing transformations for each minimal sample in the source image
        T1 = torch.zeros((matches.shape[0], 3, 3), device=dev)
        # Initialize the normalizing transformations for each minimal sample in the destination image
        T2 = torch.zeros((matches.shape[0], 3, 3), device=dev)

        # Calculate the transformation parameters
        T1[:, 0, 0] = T1[:, 1, 1] = ratio1[:]
        T2[:, 0, 0] = T2[:, 1, 1] = ratio2[:]
        T1[:, 2, 2] = T2[:, 2, 2] = 1
        T1[:, 0, 2] = -ratio1 * mass[:, 0]
        T1[:, 1, 2] = -ratio1 * mass[:, 1]
        T2[:, 2, 0] = -ratio2 * mass[:, 2]
        T2[:, 2, 1] = -ratio2 * mass[:, 3]

        return torch.cat((normalized_matches1, normalized_matches2), dim=2), T1, T2
//  [python]
def test_differentiability_jtfs_torch(random_state=42):
    device = "cpu"
    J = 8
    J_fr = 3
    shape = (4096,)
    Q = 8
    S = TimeFrequencyScatteringTorch(J=J, J_fr=J_fr, shape=shape, Q=Q, T=0, F=0).to(
        device
    )
    S.build()
    S.create_filters()
    S.load_filters()
    backend = S.backend
    torch.manual_seed(random_state)

    x = torch.randn(shape, requires_grad=True, device=device)
    x_shape = backend.shape(x)
    _, signal_shape = x_shape[:-1], x_shape[-1:]
    x_reshaped = backend.reshape_input(x, signal_shape)
    U_0_in = backend.pad(x_reshaped, pad_left=S.pad_left, pad_right=S.pad_right)

    filters = [S.phi_f, S.psi1_f, S.psi2_f]
    jtfs_gen = joint_timefrequency_scattering(
        U_0_in,
        backend,
        filters,
        S.log2_stride,
        S.average == "local",
        S.filters_fr,
        S.log2_stride_fr,
        S.average_fr == "local",
    )

    # Zeroth order
    S_0 = next(jtfs_gen)
    loss = torch.norm(S_0["coef"])
    loss.backward(retain_graph=True)
    assert torch.abs(loss) >= 0.0
    grad = x.grad
    assert torch.max(torch.abs(grad)) > 0.0

    for S in jtfs_gen:
        loss = torch.norm(S["coef"])
        loss.backward(retain_graph=True)
        assert torch.abs(loss) >= 0.0
        grad = x.grad
        assert torch.max(torch.abs(grad)) > 0.0
//  [python]
def test_differentiability_jtfs_torch(random_state=42):
    device = "cpu"
    J = 8
    J_fr = 3
    shape = (4096,)
    Q = 8
    S = TimeFrequencyScatteringTorch(J=J, J_fr=J_fr, shape=shape, Q=Q, T=0, F=0).to(
        device
    )
    S.build()
    S.create_filters()
    S.load_filters()
    backend = S.backend
    torch.manual_seed(random_state)

    x = torch.randn(shape, requires_grad=True, device=device)
    x_shape = backend.shape(x)
    _, signal_shape = x_shape[:-1], x_shape[-1:]
    x_reshaped = backend.reshape_input(x, signal_shape)
    U_0_in = backend.pad(x_reshaped, pad_left=S.pad_left, pad_right=S.pad_right)

    filters = [S.phi_f, S.psi1_f, S.psi2_f]
    jtfs_gen = joint_timefrequency_scattering(
        U_0_in,
        backend,
        filters,
        S.log2_stride,
        S.average == "local",
        S.filters_fr,
        S.log2_stride_fr,
        S.average_fr == "local",
    )

    # Zeroth order
    S_0 = next(jtfs_gen)
    loss = torch.linalg.norm(S_0["coef"])
    loss.backward(retain_graph=True)
    assert torch.abs(loss) >= 0.0
    grad = x.grad
    assert torch.max(torch.abs(grad)) > 0.0

    for S in jtfs_gen:
        loss = torch.linalg.norm(S["coef"])
        loss.backward(retain_graph=True)
        assert torch.abs(loss) >= 0.0
        grad = x.grad
        assert torch.max(torch.abs(grad)) > 0.0
//  [python]
def __init__(self, x, fun, k_easy=0.1, k_hard=0.2):

        super().__init__(x, fun)

        # When the trust-region shrinks in two consecutive
        # calculations (``tr_radius < previous_tr_radius``)
        # the lower bound ``lambda_lb`` may be reused,
        # facilitating  the convergence. To indicate no
        # previous value is known at first ``previous_tr_radius``
        # is set to -1  and ``lambda_lb`` to None.
        self.previous_tr_radius = -1
        self.lambda_lb = None

        self.niter = 0
        self.EPS = torch.finfo(x.dtype).eps

        # ``k_easy`` and ``k_hard`` are parameters used
        # to determine the stop criteria to the iterative
        # subproblem solver. Take a look at pp. 194-197
        # from reference _[1] for a more detailed description.
        self.k_easy = k_easy
        self.k_hard = k_hard

        # Get Lapack function for cholesky decomposition.
        # NOTE: cholesky_ex requires pytorch >= 1.9.0
        if 'cholesky_ex' in dir(torch.linalg):
            self.torch_cholesky = True
        else:
            # if we don't have torch cholesky, use potrf from scipy
            self.cholesky, = get_lapack_funcs(('potrf',),
                                              (self.hess.cpu().numpy(),))
            self.torch_cholesky = False

        # Get info about Hessian
        self.dimension = len(self.hess)
        self.hess_gershgorin_lb, self.hess_gershgorin_ub = gershgorin_bounds(self.hess)
        self.hess_inf = torch.norm(self.hess, float('inf'))
        self.hess_fro = torch.norm(self.hess, 'fro')

        # A constant such that for vectors smaler than that
        # backward substituition is not reliable. It was stabilished
        # based on Golub, G. H., Van Loan, C. F. (2013).
        # "Matrix computations". Forth Edition. JHU press., p.165.
        self.CLOSE_TO_ZERO = self.dimension * self.EPS * self.hess_inf
//  [python]
def __init__(self, x, fun, k_easy=0.1, k_hard=0.2):

        super().__init__(x, fun)

        # When the trust-region shrinks in two consecutive
        # calculations (``tr_radius < previous_tr_radius``)
        # the lower bound ``lambda_lb`` may be reused,
        # facilitating  the convergence. To indicate no
        # previous value is known at first ``previous_tr_radius``
        # is set to -1  and ``lambda_lb`` to None.
        self.previous_tr_radius = -1
        self.lambda_lb = None

        self.niter = 0
        self.EPS = torch.finfo(x.dtype).eps

        # ``k_easy`` and ``k_hard`` are parameters used
        # to determine the stop criteria to the iterative
        # subproblem solver. Take a look at pp. 194-197
        # from reference _[1] for a more detailed description.
        self.k_easy = k_easy
        self.k_hard = k_hard

        # Get Lapack function for cholesky decomposition.
        # NOTE: cholesky_ex requires pytorch >= 1.9.0
        if 'cholesky_ex' in dir(torch.linalg):
            self.torch_cholesky = True
        else:
            # if we don't have torch cholesky, use potrf from scipy
            self.cholesky, = get_lapack_funcs(('potrf',),
                                              (self.hess.cpu().numpy(),))
            self.torch_cholesky = False

        # Get info about Hessian
        self.dimension = len(self.hess)
        self.hess_gershgorin_lb, self.hess_gershgorin_ub = gershgorin_bounds(self.hess)
        self.hess_inf = norm(self.hess, float('inf'))
        self.hess_fro = norm(self.hess, 'fro')

        # A constant such that for vectors smaler than that
        # backward substituition is not reliable. It was stabilished
        # based on Golub, G. H., Van Loan, C. F. (2013).
        # "Matrix computations". Forth Edition. JHU press., p.165.
        self.CLOSE_TO_ZERO = self.dimension * self.EPS * self.hess_inf
//  [python]
async def aggregate_weights(self, baseline_weights, weights_received, **kwargs):
        """Perform attentive aggregation with the attention mechanism."""
        deltas_received = self.compute_weight_deltas(baseline_weights, weights_received)

        att_update = {
            name: self.trainer.zeros(weights.shape)
            for name, weights in deltas_received[0].items()
        }

        atts = OrderedDict()

        for name in baseline_weights.keys():
            atts[name] = self.trainer.zeros(len(deltas_received))
            for i, update in enumerate(deltas_received):
                # convert potential LongTensor to FloatTensor for linalg.norm
                delta = update[name].type(torch.FloatTensor)
                atts[name][i] = torch.norm(-delta)

        for name in baseline_weights.keys():
            atts[name] = F.softmax(atts[name], dim=0)

        for name, weight in baseline_weights.items():
            att_weight = self.trainer.zeros(weight.shape)
            for i, update in enumerate(deltas_received):
                delta = update[name]
                delta = delta.float()
                att_weight += torch.mul(-delta, atts[name][i])

            # Step size for aggregation used in FedAtt
            epsilon = (
                Config().algorithm.epsilon
                if hasattr(Config().algorithm, "epsilon")
                else 1.2
            )

            # The magnitude of normal noise in the randomization mechanism
            magnitude = (
                Config().algorithm.magnitude
                if hasattr(Config().algorithm, "magnitude")
                else 0.001
            )

            att_update[name] = -torch.mul(att_weight, epsilon) + torch.mul(
                torch.randn(weight.shape), magnitude
            )

        return self.update_weights(att_update)
//  [python]
async def aggregate_weights(self, baseline_weights, weights_received, **kwargs):
        """Perform attentive aggregation with the attention mechanism."""
        deltas_received = self.compute_weight_deltas(baseline_weights, weights_received)

        att_update = {
            name: self.trainer.zeros(weights.shape)
            for name, weights in deltas_received[0].items()
        }

        atts = OrderedDict()

        for name in baseline_weights.keys():
            atts[name] = self.trainer.zeros(len(deltas_received))
            for i, update in enumerate(deltas_received):
                # convert potential LongTensor to FloatTensor for linalg.norm
                delta = update[name].type(torch.FloatTensor)
                atts[name][i] = torch.linalg.norm(-delta)

        for name in baseline_weights.keys():
            atts[name] = F.softmax(atts[name], dim=0)

        for name, weight in baseline_weights.items():
            att_weight = self.trainer.zeros(weight.shape)
            for i, update in enumerate(deltas_received):
                delta = update[name]
                delta = delta.float()
                att_weight += torch.mul(-delta, atts[name][i])

            # Step size for aggregation used in FedAtt
            epsilon = (
                Config().algorithm.epsilon
                if hasattr(Config().algorithm, "epsilon")
                else 1.2
            )

            # The magnitude of normal noise in the randomization mechanism
            magnitude = (
                Config().algorithm.magnitude
                if hasattr(Config().algorithm, "magnitude")
                else 0.001
            )

            att_update[name] = -torch.mul(att_weight, epsilon) + torch.mul(
                torch.randn(weight.shape), magnitude
            )

        return self.update_weights(att_update)
//  [python]
def forward(self, batch):

        # Edge features
        vectors = batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]]  # [n_edges, 3]
        lengths = torch.norm(vectors, dim=-1, keepdim=True)  # [n_edges, 1]
        
        h_V = self.emb_in(batch.atoms)  # (n,) -> (n, d)
        h_E = (
            self.radial_embedding(lengths), 
            torch.nan_to_num(torch.div(vectors, lengths)).unsqueeze_(-2)
        )

        h_V = self.W_v(h_V)
        h_E = self.W_e(h_E)
    
        for layer in self.layers:
            h_V = layer(h_V, batch.edge_index, h_E)

        out = self.pool(gvp._merge(*h_V), batch.batch)  # (n, d) -> (batch_size, d)
        
        if not self.equivariant_pred:
            # Select only scalars for invariant prediction
            out = out[:,:self.s_dim]
        
        return self.pred(out)
//  [python]
def forward(self, batch):

        # Edge features
        vectors = batch.pos[batch.edge_index[0]] - batch.pos[batch.edge_index[1]]  # [n_edges, 3]
        lengths = torch.linalg.norm(vectors, dim=-1, keepdim=True)  # [n_edges, 1]
        
        h_V = self.emb_in(batch.atoms)  # (n,) -> (n, d)
        h_E = (
            self.radial_embedding(lengths), 
            torch.nan_to_num(torch.div(vectors, lengths)).unsqueeze_(-2)
        )

        h_V = self.W_v(h_V)
        h_E = self.W_e(h_E)
    
        for layer in self.layers:
            h_V = layer(h_V, batch.edge_index, h_E)

        out = self.pool(gvp._merge(*h_V), batch.batch)  # (n, d) -> (batch_size, d)
        
        if not self.equivariant_pred:
            # Select only scalars for invariant prediction
            out = out[:,:self.s_dim]
        
        return self.pred(out)
//  [python]
def _get_scale_all(self, coords: torch.Tensor, eps: float=1e-12) -> torch.Tensor:
        a_x1 = torch.norm(coords[:, 0, :] - coords[:, 1, :], p=2, dim=1)
        a_y1 = torch.norm(coords[:, 1, :] - coords[:, 2, :], p=2, dim=1)
        a_x2 = torch.norm(coords[:, 2, :] - coords[:, 3, :], p=2, dim=1)
        a_y2 = torch.norm(coords[:, 3, :] - coords[:, 0, :], p=2, dim=1)
        a_x = 0.5 * (a_x1 + a_x2)
        a_y = 0.5 * (a_y1 + a_y2)
        return torch.stack([a_y, a_x], dim=1)
//  [python]
def _get_scale_all(self, coords: torch.Tensor, eps: float=1e-12) -> torch.Tensor:
        a_x1 = torch.linalg.norm(coords[:, 0, :] - coords[:, 1, :], ord=2, dim=1)
        a_y1 = torch.linalg.norm(coords[:, 1, :] - coords[:, 2, :], ord=2, dim=1)
        a_x2 = torch.linalg.norm(coords[:, 2, :] - coords[:, 3, :], ord=2, dim=1)
        a_y2 = torch.linalg.norm(coords[:, 3, :] - coords[:, 0, :], ord=2, dim=1)
        a_x = 0.5 * (a_x1 + a_x2)
        a_y = 0.5 * (a_y1 + a_y2)
        return torch.stack([a_y, a_x], dim=1)
//  [python]
def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        forward function

        Parameters:
        -----------
        input: torch.Tensor
            contains the sampled electrons coordinate in the shape (batch_size,number_of_electrons*3)

        Returns:
        --------
        psi: torch.Tensor
            contains the wavefunction - 'psi' value. It is in the shape (batch_size), where each row corresponds to the solution of one of the batches
        """
        # creating one and two electron features
        eps = torch.tensor(1e-36)
        self.input = input.reshape((self.batch_size, -1, 3))
        two_electron_vector = self.input.unsqueeze(1) - self.input.unsqueeze(2)
        two_electron_distance = torch.norm(two_electron_vector + eps,
                                                  dim=3).unsqueeze(3)
        two_electron = torch.cat((two_electron_vector, two_electron_distance),
                                 dim=3)
        two_electron = torch.reshape(
            two_electron,
            (self.batch_size, self.total_electron, self.total_electron, -1))
        one_electron_vector = self.input.unsqueeze(
            1) - self.nucleon_pos.unsqueeze(1)
        one_electron_distance = torch.norm(one_electron_vector, dim=3)
        one_electron = torch.cat(
            (one_electron_vector, one_electron_distance.unsqueeze(-1)), dim=3)
        one_electron = torch.reshape(one_electron.permute(0, 2, 1, 3),
                                     (self.batch_size, self.total_electron, -1))
        one_electron_vector_permuted = one_electron_vector.permute(0, 2, 1,
                                                                   3).float()
        # setting the fermient layer and fermient envelope layer batch size to be that of the current batch size of the model. This enables for vectorized calculations of hessians and jacobians.
        self.ferminet_layer[0].batch_size = self.batch_size
        self.ferminet_layer_envelope[0].batch_size = self.batch_size
        one_electron, _ = self.ferminet_layer[0].forward(
            one_electron.to(torch.float32), two_electron.to(torch.float32))
        self.psi, self.psi_up, self.psi_down = self.ferminet_layer_envelope[
            0].forward(one_electron, one_electron_vector_permuted)
        return self.psi
//  [python]
def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        forward function

        Parameters:
        -----------
        input: torch.Tensor
            contains the sampled electrons coordinate in the shape (batch_size,number_of_electrons*3)

        Returns:
        --------
        psi: torch.Tensor
            contains the wavefunction - 'psi' value. It is in the shape (batch_size), where each row corresponds to the solution of one of the batches
        """
        # creating one and two electron features
        eps = torch.tensor(1e-36)
        self.input = input.reshape((self.batch_size, -1, 3))
        two_electron_vector = self.input.unsqueeze(1) - self.input.unsqueeze(2)
        two_electron_distance = torch.linalg.norm(two_electron_vector + eps,
                                                  dim=3).unsqueeze(3)
        two_electron = torch.cat((two_electron_vector, two_electron_distance),
                                 dim=3)
        two_electron = torch.reshape(
            two_electron,
            (self.batch_size, self.total_electron, self.total_electron, -1))
        one_electron_vector = self.input.unsqueeze(
            1) - self.nucleon_pos.unsqueeze(1)
        one_electron_distance = torch.linalg.norm(one_electron_vector, dim=3)
        one_electron = torch.cat(
            (one_electron_vector, one_electron_distance.unsqueeze(-1)), dim=3)
        one_electron = torch.reshape(one_electron.permute(0, 2, 1, 3),
                                     (self.batch_size, self.total_electron, -1))
        one_electron_vector_permuted = one_electron_vector.permute(0, 2, 1,
                                                                   3).float()
        # setting the fermient layer and fermient envelope layer batch size to be that of the current batch size of the model. This enables for vectorized calculations of hessians and jacobians.
        self.ferminet_layer[0].batch_size = self.batch_size
        self.ferminet_layer_envelope[0].batch_size = self.batch_size
        one_electron, _ = self.ferminet_layer[0].forward(
            one_electron.to(torch.float32), two_electron.to(torch.float32))
        self.psi, self.psi_up, self.psi_down = self.ferminet_layer_envelope[
            0].forward(one_electron, one_electron_vector_permuted)
        return self.psi
//  [python]
def _project_points(self,
        points3d:           torch.Tensor,
        epsilon:            float=1e-12,
    ):
        phi = torch.atan2(points3d[:, :, 0], -1.0 * points3d[:, :, 1] + epsilon) # [-pi, pi]
        xy_dist = torch.norm(points3d[:, :, :2], dim=-1)
        theta = -1.0 * torch.atan2(points3d[:, :, 2], xy_dist + epsilon)  # [-pi / 2.0, pi / 2.0]
        u = phi / np.pi
        v = theta / (0.5 * np.pi)
        return torch.stack([u, v], dim=-1)
//  [python]
def _project_points(self,
        points3d:           torch.Tensor,
        epsilon:            float=1e-12,
    ):
        phi = torch.atan2(points3d[:, :, 0], -1.0 * points3d[:, :, 1] + epsilon) # [-pi, pi]
        xy_dist = torch.linalg.norm(points3d[:, :, :2], ord=2, dim=-1)
        theta = -1.0 * torch.atan2(points3d[:, :, 2], xy_dist + epsilon)  # [-pi / 2.0, pi / 2.0]
        u = phi / np.pi
        v = theta / (0.5 * np.pi)
        return torch.stack([u, v], dim=-1)
//  [python]
def loss(self):
        def f(z):
            return (1.0 / self.A.size()[0]) * torch.norm(
                torch.max(self.A.mm(z.view(self.βs.size())), dim=1)[0] - self.y, 1
            )

        return f
//  [python]
def loss(self):
        def f(z):
            return (1.0 / self.A.size()[0]) * torch.linalg.norm(
                torch.max(self.A.mm(z.view(self.βs.size())), dim=1)[0] - self.y, 1
            )

        return f
//  [python]
def _homography_floor_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
    ):
        b, N, _ = top_corners.size()
        u = bottom_corners[:, :, 0] * np.pi
        v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        c = floor_z / torch.tan(v)
        x = c * torch.sin(u)
        y = -c * torch.cos(u)
        floor_xy = torch.stack([x, y], dim=-1)
        scale = self._get_scale_all(floor_xy)
        scale = scale / 2.0
        centroid = floor_xy.mean(dim=1)
        c = torch.norm(floor_xy, p=2, dim=-1)
        v = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
        ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        floor_xy = floor_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(floor_xy[..., 0], floor_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(floor_xy, axes)    
        homogeneous = torch.cat([floor_xy, torch.ones_like(floor_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        original_xy = floor_xy + centroid.unsqueeze(1)
        R, t, s = self._svd(rect_floor_xy, original_xy[:, inds.squeeze(), :])
        rect_floor_xy = self._transform_points(rect_floor_xy, R, t, s)
        bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
        return top_points, bottom_points
//  [python]
def _homography_floor_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        floor_z: float=-1.6,
    ):
        b, N, _ = top_corners.size()
        u = bottom_corners[:, :, 0] * np.pi
        v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        c = floor_z / torch.tan(v)
        x = c * torch.sin(u)
        y = -c * torch.cos(u)
        floor_xy = torch.stack([x, y], dim=-1)
        scale = self._get_scale_all(floor_xy)
        scale = scale / 2.0
        centroid = floor_xy.mean(dim=1)
        c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
        v = top_corners[:, :, 1] * (-0.5 * np.pi)
        ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
        ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
        floor_xy = floor_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(floor_xy[..., 0], floor_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(floor_xy, axes)    
        homogeneous = torch.cat([floor_xy, torch.ones_like(floor_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        original_xy = floor_xy + centroid.unsqueeze(1)
        R, t, s = self._svd(rect_floor_xy, original_xy[:, inds.squeeze(), :])
        rect_floor_xy = self._transform_points(rect_floor_xy, R, t, s)
        bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
        return top_points, bottom_points
//  [python]
def dist_sparse(self, x: Tensor):
        d = self.x.size(0)
        return torch.norm(
            torch.topk(torch.abs(x), d - self.k, largest=False).values,
        )
//  [python]
def dist_sparse(self, x: Tensor):
        d = self.x.size(0)
        return torch.linalg.norm(
            torch.topk(torch.abs(x), d - self.k, largest=False).values,
        )
//  [python]
def loss_altproj(self):
        m = self.y.size(-1)

        def f(z):
            z_comp = z[:m] + z[m:] * 1j
            return torch.norm(
                z_comp - self.A @ (torch.linalg.lstsq(self.A, z_comp).solution)
            ) + torch.norm(z_comp - self.y * phase(z_comp))

        return f
//  [python]
def loss_altproj(self):
        m = self.y.size(-1)

        def f(z):
            z_comp = z[:m] + z[m:] * 1j
            return torch.linalg.norm(
                z_comp - self.A @ (torch.linalg.lstsq(self.A, z_comp).solution)
            ) + torch.linalg.norm(z_comp - self.y * phase(z_comp))

        return f
//  [python]
def _homography_ceil_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        u_t = top_corners[:, :, 0] * np.pi
        v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        c = ceil_z / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        c = torch.norm(ceil_xy, p=2, dim=-1)
        v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_z = (c * torch.tan(v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / floor_z
        floor_z = -ceil_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        c = ceil_z_fix.squeeze(-1) / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        scale = self._get_scale_all(ceil_xy)
        scale = scale / 2.0
        centroid = ceil_xy.mean(dim=1)
        ceil_xy = ceil_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(ceil_xy[..., 0], ceil_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(ceil_xy, axes)
        homogeneous = torch.cat([ceil_xy, torch.ones_like(ceil_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_ceil_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        original_xy = ceil_xy + centroid.unsqueeze(1)
        R, t, s = self._svd(rect_ceil_xy, original_xy[:, inds.squeeze(), :])
        rect_ceil_xy = self._transform_points(rect_ceil_xy, R, t, s)
        bottom_points = torch.cat([rect_ceil_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_ceil_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points
//  [python]
def _homography_ceil_svd(self,
        top_corners: torch.Tensor, # in [-1, 1]
        bottom_corners: torch.Tensor, # in [-1, 1]
        ceil_z: float=1.6,
    ):
        b, N, _ = top_corners.size()
        u_t = top_corners[:, :, 0] * np.pi
        v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
        c = ceil_z / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
        v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
        floor_z = (c * torch.tan(v_b)).mean(dim=1, keepdim=True)
        fix_ceil = -ceil_z / floor_z
        floor_z = -ceil_z
        ceil_z_fix = ceil_z * fix_ceil
        ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
        c = ceil_z_fix.squeeze(-1) / torch.tan(v_t)
        x = c * torch.sin(u_t)
        y = -c * torch.cos(u_t)
        ceil_xy = torch.stack([x, y], dim=-1)
        scale = self._get_scale_all(ceil_xy)
        scale = scale / 2.0
        centroid = ceil_xy.mean(dim=1)
        ceil_xy = ceil_xy - centroid.unsqueeze(1)
        inds = torch.sort(torch.atan2(ceil_xy[..., 0], ceil_xy[..., 1] + 1e-12))[1]
        axes = self.cuboid_axes[:, inds.squeeze(), :]
        homography = kornia.get_perspective_transform(ceil_xy, axes)
        homogeneous = torch.cat([ceil_xy, torch.ones_like(ceil_xy[..., -1:])], dim=2)
        xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
        xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
        rect_ceil_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
        original_xy = ceil_xy + centroid.unsqueeze(1)
        R, t, s = self._svd(rect_ceil_xy, original_xy[:, inds.squeeze(), :])
        rect_ceil_xy = self._transform_points(rect_ceil_xy, R, t, s)
        bottom_points = torch.cat([rect_ceil_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
        top_points = torch.cat([rect_ceil_xy, ceil_z_fix], dim=-1)
        return top_points, bottom_points
//  [python]
def __getitem__(self, idxs: List[int]) -> Dict[str, object]:
        frame_dicts = []
        for i in idxs:
            sample = self.samples[i]
            frame_dict = self._getitem(sample)
            frame_dicts.append(frame_dict)

        dists = [torch.tensor(0)]
        for i in range(1, len(frame_dicts)):
            # Get the relative distance from the prev sample to the curr sample
            curr_fd = frame_dicts[i]
            prev_fd = frame_dicts[i - 1]
            dist = torch.norm(curr_fd["cam_T"][:3, 3] - prev_fd["cam_T"][:3, 3])
            dists.append(dist)

        cam_Ts = [fd["cam_T"] for fd in frame_dicts]

        frame_Ts = [torch.eye(4)]
        for i in range(1, len(frame_dicts)):
            # Get the relative frame_T for this frame by comparing the current cam_T to the previous cam_T
            frame_T = torch.linalg.solve(cam_Ts[i - 1], cam_Ts[i])
            # TODO use torch.solve instead
            # frame_T = torch.linalg.solve(cam_Ts[i], cam_Ts[i - 1])[0]
            frame_Ts.append(frame_T)

        imgs = [fd["color"] for fd in frame_dicts]
        token = [fd["token"] for fd in frame_dicts]

        # timestamp is in microseconds, need to convert it to seconds and
        # normalize to first frame
        frame_time = torch.tensor(
            [fd["frame_time"] for fd in frame_dicts], dtype=torch.int64
        )

        long_cam_Ts = [
            get_ego_T(self.nusc, sample_data) for sample_data in self.samples[idxs[0] :]
        ]

        mask_transform = transforms.Compose(
            [
                transforms.PILToTensor(),
                transforms.ConvertImageDtype(
                    torch.bfloat16
                ),  # drop precision to save memory
            ]
        )
        mask_path = os.path.join(
            os.path.dirname(__file__),
            "nuscenes_masks",
            f"{self.sensor}.png",
        )
        mask = Image.open(mask_path).convert("L")
        mask = mask_transform(mask)
        assert mask.shape == (1, 480, 640), mask.shape

        return {
            "weight": torch.tensor(frame_dicts[0]["weight"]),
            "distance": torch.stack(dists),
            "cam_T": torch.stack(cam_Ts),
            "frame_T": torch.stack(frame_Ts),
            "long_cam_T": torch.stack(long_cam_Ts),
            "frame_time": frame_time,
            "K": frame_dicts[0][
                "K"
            ],  # only one cam to pixel matrix is required as it doesn't change during drive
            "T": frame_dicts[0][
                "T"
            ],  # only one cam to car translation matrix is required as it doesn't change during drive
            "color": torch.stack(imgs),
            "mask": mask,
            "token": token,
        }
//  [python]
def __getitem__(self, idxs: List[int]) -> Dict[str, object]:
        frame_dicts = []
        for i in idxs:
            sample = self.samples[i]
            frame_dict = self._getitem(sample)
            frame_dicts.append(frame_dict)

        dists = [torch.tensor(0)]
        for i in range(1, len(frame_dicts)):
            # Get the relative distance from the prev sample to the curr sample
            curr_fd = frame_dicts[i]
            prev_fd = frame_dicts[i - 1]
            dist = torch.linalg.norm(curr_fd["cam_T"][:3, 3] - prev_fd["cam_T"][:3, 3])
            dists.append(dist)

        cam_Ts = [fd["cam_T"] for fd in frame_dicts]

        frame_Ts = [torch.eye(4)]
        for i in range(1, len(frame_dicts)):
            # Get the relative frame_T for this frame by comparing the current cam_T to the previous cam_T
            frame_T = torch.linalg.solve(cam_Ts[i - 1], cam_Ts[i])
            # TODO use torch.solve instead
            # frame_T = torch.linalg.solve(cam_Ts[i], cam_Ts[i - 1])[0]
            frame_Ts.append(frame_T)

        imgs = [fd["color"] for fd in frame_dicts]
        token = [fd["token"] for fd in frame_dicts]

        # timestamp is in microseconds, need to convert it to seconds and
        # normalize to first frame
        frame_time = torch.tensor(
            [fd["frame_time"] for fd in frame_dicts], dtype=torch.int64
        )

        long_cam_Ts = [
            get_ego_T(self.nusc, sample_data) for sample_data in self.samples[idxs[0] :]
        ]

        mask_transform = transforms.Compose(
            [
                transforms.PILToTensor(),
                transforms.ConvertImageDtype(
                    torch.bfloat16
                ),  # drop precision to save memory
            ]
        )
        mask_path = os.path.join(
            os.path.dirname(__file__),
            "nuscenes_masks",
            f"{self.sensor}.png",
        )
        mask = Image.open(mask_path).convert("L")
        mask = mask_transform(mask)
        assert mask.shape == (1, 480, 640), mask.shape

        return {
            "weight": torch.tensor(frame_dicts[0]["weight"]),
            "distance": torch.stack(dists),
            "cam_T": torch.stack(cam_Ts),
            "frame_T": torch.stack(frame_Ts),
            "long_cam_T": torch.stack(long_cam_Ts),
            "frame_time": frame_time,
            "K": frame_dicts[0][
                "K"
            ],  # only one cam to pixel matrix is required as it doesn't change during drive
            "T": frame_dicts[0][
                "T"
            ],  # only one cam to car translation matrix is required as it doesn't change during drive
            "color": torch.stack(imgs),
            "mask": mask,
            "token": token,
        }
//  [python]
def R_to_allocentric(K, R, u=None, v=None):
    """
    Convert a rotation matrix or series of rotation matrices to allocentric
    representation given a 2D location (u, v) in pixels. 
    When u or v are not available, we fall back on the principal point of K.
    """
    if type(K) == torch.Tensor:
        fx = K[:, 0, 0]
        fy = K[:, 1, 1]
        sx = K[:, 0, 2]
        sy = K[:, 1, 2]

        n = len(K)
        
        oray = torch.stack(((u - sx)/fx, (v - sy)/fy, torch.ones_like(u))).T
        oray = oray / torch.norm(oray, dim=1).unsqueeze(1)
        angle = torch.acos(oray[:, -1])

        axis = torch.zeros_like(oray)
        axis[:, 0] = axis[:, 0] - oray[:, 1]
        axis[:, 1] = axis[:, 1] + oray[:, 0]
        norms = torch.norm(axis, dim=1)

        valid_angle = angle > 0

        M = axis_angle_to_matrix(angle.unsqueeze(1)*axis/norms.unsqueeze(1))
        
        R_view = R.clone()
        R_view[valid_angle] = torch.bmm(M[valid_angle].transpose(2, 1), R[valid_angle])

    else:
        fx = K[0][0]
        fy = K[1][1]
        sx = K[0][2]
        sy = K[1][2]
        
        if u is None:
            u = sx

        if v is None:
            v = sy

        oray = np.array([(u - sx)/fx, (v - sy)/fy, 1])
        oray = oray / np.linalg.norm(oray)
        cray = np.array([0, 0, 1])
        angle = math.acos(cray.dot(oray))
        if angle != 0:
            axis = np.cross(cray, oray)
            axis_torch = torch.from_numpy(angle*axis/np.linalg.norm(axis)).float()
            R_view = np.dot(axis_angle_to_matrix(axis_torch).numpy().T, R)
        else: 
            R_view = R

    return R_view
//  [python]
def R_to_allocentric(K, R, u=None, v=None):
    """
    Convert a rotation matrix or series of rotation matrices to allocentric
    representation given a 2D location (u, v) in pixels. 
    When u or v are not available, we fall back on the principal point of K.
    """
    if type(K) == torch.Tensor:
        fx = K[:, 0, 0]
        fy = K[:, 1, 1]
        sx = K[:, 0, 2]
        sy = K[:, 1, 2]

        n = len(K)
        
        oray = torch.stack(((u - sx)/fx, (v - sy)/fy, torch.ones_like(u))).T
        oray = oray / torch.linalg.norm(oray, dim=1).unsqueeze(1)
        angle = torch.acos(oray[:, -1])

        axis = torch.zeros_like(oray)
        axis[:, 0] = axis[:, 0] - oray[:, 1]
        axis[:, 1] = axis[:, 1] + oray[:, 0]
        norms = torch.linalg.norm(axis, dim=1)

        valid_angle = angle > 0

        M = axis_angle_to_matrix(angle.unsqueeze(1)*axis/norms.unsqueeze(1))
        
        R_view = R.clone()
        R_view[valid_angle] = torch.bmm(M[valid_angle].transpose(2, 1), R[valid_angle])

    else:
        fx = K[0][0]
        fy = K[1][1]
        sx = K[0][2]
        sy = K[1][2]
        
        if u is None:
            u = sx

        if v is None:
            v = sy

        oray = np.array([(u - sx)/fx, (v - sy)/fy, 1])
        oray = oray / np.linalg.norm(oray)
        cray = np.array([0, 0, 1])
        angle = math.acos(cray.dot(oray))
        if angle != 0:
            axis = np.cross(cray, oray)
            axis_torch = torch.from_numpy(angle*axis/np.linalg.norm(axis)).float()
            R_view = np.dot(axis_angle_to_matrix(axis_torch).numpy().T, R)
        else: 
            R_view = R

    return R_view
//  [python]
def R_from_allocentric(K, R_view, u=None, v=None):
    """
    Convert a rotation matrix or series of rotation matrices to egocentric
    representation given a 2D location (u, v) in pixels. 
    When u or v are not available, we fall back on the principal point of K.
    """
    if type(K) == torch.Tensor:
        fx = K[:, 0, 0]
        fy = K[:, 1, 1]
        sx = K[:, 0, 2]
        sy = K[:, 1, 2]

        n = len(K)
        
        oray = torch.stack(((u - sx)/fx, (v - sy)/fy, torch.ones_like(u))).T
        oray = oray / torch.norm(oray, dim=1).unsqueeze(1)
        angle = torch.acos(oray[:, -1])

        axis = torch.zeros_like(oray)
        axis[:, 0] = axis[:, 0] - oray[:, 1]
        axis[:, 1] = axis[:, 1] + oray[:, 0]
        norms = torch.norm(axis, dim=1)

        valid_angle = angle > 0

        M = axis_angle_to_matrix(angle.unsqueeze(1)*axis/norms.unsqueeze(1))
        
        R = R_view.clone()
        R[valid_angle] = torch.bmm(M[valid_angle], R_view[valid_angle])

    else:
        fx = K[0][0]
        fy = K[1][1]
        sx = K[0][2]
        sy = K[1][2]
        
        if u is None:
            u = sx

        if v is None:
            v = sy

        oray = np.array([(u - sx)/fx, (v - sy)/fy, 1])
        oray = oray / np.linalg.norm(oray)
        cray = np.array([0, 0, 1])
        angle = math.acos(cray.dot(oray))
        if angle != 0:
            #axis = np.cross(cray, oray)
            axis = np.array([-oray[1], oray[0], 0])
            axis_torch = torch.from_numpy(angle*axis/np.linalg.norm(axis)).float()
            R = np.dot(axis_angle_to_matrix(axis_torch).numpy(), R_view)
        else: 
            R = R_view

    return R
//  [python]
def R_from_allocentric(K, R_view, u=None, v=None):
    """
    Convert a rotation matrix or series of rotation matrices to egocentric
    representation given a 2D location (u, v) in pixels. 
    When u or v are not available, we fall back on the principal point of K.
    """
    if type(K) == torch.Tensor:
        fx = K[:, 0, 0]
        fy = K[:, 1, 1]
        sx = K[:, 0, 2]
        sy = K[:, 1, 2]

        n = len(K)
        
        oray = torch.stack(((u - sx)/fx, (v - sy)/fy, torch.ones_like(u))).T
        oray = oray / torch.linalg.norm(oray, dim=1).unsqueeze(1)
        angle = torch.acos(oray[:, -1])

        axis = torch.zeros_like(oray)
        axis[:, 0] = axis[:, 0] - oray[:, 1]
        axis[:, 1] = axis[:, 1] + oray[:, 0]
        norms = torch.linalg.norm(axis, dim=1)

        valid_angle = angle > 0

        M = axis_angle_to_matrix(angle.unsqueeze(1)*axis/norms.unsqueeze(1))
        
        R = R_view.clone()
        R[valid_angle] = torch.bmm(M[valid_angle], R_view[valid_angle])

    else:
        fx = K[0][0]
        fy = K[1][1]
        sx = K[0][2]
        sy = K[1][2]
        
        if u is None:
            u = sx

        if v is None:
            v = sy

        oray = np.array([(u - sx)/fx, (v - sy)/fy, 1])
        oray = oray / np.linalg.norm(oray)
        cray = np.array([0, 0, 1])
        angle = math.acos(cray.dot(oray))
        if angle != 0:
            #axis = np.cross(cray, oray)
            axis = np.array([-oray[1], oray[0], 0])
            axis_torch = torch.from_numpy(angle*axis/np.linalg.norm(axis)).float()
            R = np.dot(axis_angle_to_matrix(axis_torch).numpy(), R_view)
        else: 
            R = R_view

    return R
//  [python]
def loss(self):
        m = self.y.size(0)
        return lambda Z: (1 / m) * torch.norm(
            self.y - torch.sum(self.L.mm(Z) * self.R.mm(Z), dim=1), 1
        )
//  [python]
def loss(self):
        m = self.y.size(0)
        return lambda Z: (1 / m) * torch.linalg.norm(
            self.y - torch.sum(self.L.mm(Z) * self.R.mm(Z), dim=1), 1
        )
//  [python]
def initializer(self, δ):
        Δ = torch.randn(self.X.size(), dtype=torch.double)
        Δ = Δ / torch.norm(Δ)
        return (self.X + δ * Δ).requires_grad_(True)