jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3, qef_reg_scale):
        p_bxnx3 = p_bxnx3.reshape(-1, 7, 3)
        norm_bxnx3 = norm_bxnx3.reshape(-1, 7, 3)
        c_bx3 = c_bx3.reshape(-1, 3)
        A = norm_bxnx3
        B = ((p_bxnx3) * norm_bxnx3).sum(-1, keepdims=True)

        A_reg = (torch.eye(3, device=p_bxnx3.device) * qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
        B_reg = (qef_reg_scale * c_bx3).unsqueeze(-1)
        A = torch.cat([A, A_reg], 1)
        B = torch.cat([B, B_reg], 1)
        dual_verts = torch.solve(B, A.transpose(-2,-1)).solution.squeeze(-1)
        return dual_verts
//  [python]
def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3, qef_reg_scale):
        p_bxnx3 = p_bxnx3.reshape(-1, 7, 3)
        norm_bxnx3 = norm_bxnx3.reshape(-1, 7, 3)
        c_bx3 = c_bx3.reshape(-1, 3)
        A = norm_bxnx3
        B = ((p_bxnx3) * norm_bxnx3).sum(-1, keepdims=True)

        A_reg = (torch.eye(3, device=p_bxnx3.device) * qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
        B_reg = (qef_reg_scale * c_bx3).unsqueeze(-1)
        A = torch.cat([A, A_reg], 1)
        B = torch.cat([B, B_reg], 1)
        dual_verts = torch.linalg.lstsq(A, B).solution.squeeze(-1)
        return dual_verts
//  [python]
def Kron_reduction_translational(edge_index,restriction_maps,boundary_vertices,interior_vertices):
    L = Laplacian(edge_index, restriction_maps)
    d = coboundary(edge_index, restriction_maps)
    dv = restriction_maps.shape[-1]
    nbatch = L.shape[0]

    Bidx = get_matrix_indices(boundary_vertices,dv)
    Uidx = get_matrix_indices(interior_vertices,dv)
    batch_idx = np.arange(nbatch)

    LBB = L[np.ix_(batch_idx,Bidx,Bidx)]
    LUB = L[np.ix_(batch_idx,Uidx,Bidx)]
    LUU = L[np.ix_(batch_idx,Uidx,Uidx)]

    didx = np.arange(d.shape[1])
    dU = d[np.ix_(batch_idx, didx, Uidx)]
    dB = d[np.ix_(batch_idx, didx, Bidx)]

    invLUB = torch.solve(LUB,LUU).solution

    schur = LBB - torch.transpose(LUB, 1, 2) @ invLUB
    affine = -dU @ invLUB + dB
    return schur, affine
//  [python]
def Kron_reduction_translational(edge_index,restriction_maps,boundary_vertices,interior_vertices):
    L = Laplacian(edge_index, restriction_maps)
    d = coboundary(edge_index, restriction_maps)
    dv = restriction_maps.shape[-1]
    nbatch = L.shape[0]

    Bidx = get_matrix_indices(boundary_vertices,dv)
    Uidx = get_matrix_indices(interior_vertices,dv)
    batch_idx = np.arange(nbatch)

    LBB = L[np.ix_(batch_idx,Bidx,Bidx)]
    LUB = L[np.ix_(batch_idx,Uidx,Bidx)]
    LUU = L[np.ix_(batch_idx,Uidx,Uidx)]

    didx = np.arange(d.shape[1])
    dU = d[np.ix_(batch_idx, didx, Uidx)]
    dB = d[np.ix_(batch_idx, didx, Bidx)]

    invLUB = torch.linalg.lstsq(LUU, LUB).solution

    schur = LBB - torch.transpose(LUB, 1, 2) @ invLUB
    affine = -dU @ invLUB + dB
    return schur, affine
//  [python]
def get_targets(xs, ys, target_mode):
    targets = []
    n_points = xs.shape[1]
    n_dims = xs.shape[2]
    if target_mode == 'grad':
        for i_n in range(1, n_points + 1):
            target = (ys[:, None, :i_n] @ xs[:, :i_n, :])[:, 0, :] / n_points  # [B, 1, d] -> [B, d]
            targets.append(target)
    elif target_mode == 'Wols':
        device = xs.device
        xs, ys = xs.cpu(), ys.cpu()
        for i_n in range(1, n_dims + 1):
            train_xs, train_ys = xs[:, :i_n], ys[:, :i_n]
            target, _, _, _ = torch.solve(train_ys.unsqueeze(2), train_xs.t())  # [B, d, 1]
            target = target[:, :, 0]  # -> [B, d]
            targets.append(target.to(device))
    else:
        raise NotImplementedError

    return targets
//  [python]
def get_targets(xs, ys, target_mode):
    targets = []
    n_points = xs.shape[1]
    n_dims = xs.shape[2]
    if target_mode == 'grad':
        for i_n in range(1, n_points + 1):
            target = (ys[:, None, :i_n] @ xs[:, :i_n, :])[:, 0, :] / n_points  # [B, 1, d] -> [B, d]
            targets.append(target)
    elif target_mode == 'Wols':
        device = xs.device
        xs, ys = xs.cpu(), ys.cpu()
        for i_n in range(1, n_dims + 1):
            train_xs, train_ys = xs[:, :i_n], ys[:, :i_n]
            target, _, _, _ = torch.linalg.lstsq(train_xs, train_ys.unsqueeze(2))  # [B, d, 1]
            target = target[:, :, 0]  # -> [B, d]
            targets.append(target.to(device))
    else:
        raise NotImplementedError

    return targets
//  [python]
def _lstsq(A: torch.Tensor, B: torch.Tensor) ->torch.Tensor:
    if A.shape[0] != B.shape[0]:
        raise RuntimeError(f'A has different number of rows than B! A.shape = {A.shape}, B.shape = {B.shape}')
    if A.ndim != 2:
        raise RuntimeError('A is not a rank 2 tensor: has shape', A.shape)
    if B.ndim != 2:
        raise RuntimeError('B is not a rank 2 tensor: has shape', A.shape)
    return torch.solve(torch.matmul(A.T,A), torch.matmul(A.T,B)).solution
//  [python]
def _lstsq(A: torch.Tensor, B: torch.Tensor) ->torch.Tensor:
    if A.shape[0] != B.shape[0]:
        raise RuntimeError(f'A has different number of rows than B! A.shape = {A.shape}, B.shape = {B.shape}')
    if A.ndim != 2:
        raise RuntimeError('A is not a rank 2 tensor: has shape', A.shape)
    if B.ndim != 2:
        raise RuntimeError('B is not a rank 2 tensor: has shape', A.shape)
    return torch.linalg.lstsq(A, B).solution
//  [python]
def _get_nu(self, fY, hY):
        """Compute nu (ie lambda) if not provided by the problem's solver.
        That is, solve: hY^T nu = fY^T.
        """
        p = hY.size(1)
        nu = fY.new_zeros(self.b, p)
        for i in range(self.b): # loop over batch
            solution = torch.lstsq(fY[i, :].unsqueeze(-1), hY[i, :, :].t())[0]
            nu[i, :] = solution[:p, :].squeeze() # extract first p values
        return nu
//  [python]
def _get_nu(self, fY, hY):
        """Compute nu (ie lambda) if not provided by the problem's solver.
        That is, solve: hY^T nu = fY^T.
        """
        p = hY.size(1)
        nu = fY.new_zeros(self.b, p)
        for i in range(self.b): # loop over batch
            solution = torch.linalg.lstsq(hY[i, :, :].t(), fY[i, :].unsqueeze(-1))[0]
            nu[i, :] = solution[:p, :].squeeze() # extract first p values
        return nu
//  [python]
def f(lamb: float) -> torch.Tensor:
    """f.

    :param lamb:
    :type lamb: float
    :rtype: torch.Tensor
    """
    y = y_true + 0.3 * torch.randn(25)

    lhs = X.T @ X
    rhs = X.T @ y
    ridge = lamb * torch.eye(lhs.shape[0])
    w = torch.solve(rhs, lhs + ridge)[0]
    y_hat = X @ w
    return y_hat
//  [python]
def f(lamb: float) -> torch.Tensor:
    """f.

    :param lamb:
    :type lamb: float
    :rtype: torch.Tensor
    """
    y = y_true + 0.3 * torch.randn(25)

    lhs = X.T @ X
    rhs = X.T @ y
    ridge = lamb * torch.eye(lhs.shape[0])
    w = torch.linalg.lstsq(lhs + ridge, rhs).solution
    y_hat = X @ w
    return y_hat
//  [python]
def gauss_newton_step(
        p: torch.Tensor,
        function: Callable,
        jac_function: Callable,
        wvec: torch.Tensor,
        scon: torch.Tensor = None,
        l: float = 1.,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Gauss-Newton step function for parallel least-squares fitting of non-linear functions

    :param p: current guess
    :param function: user-provided function which takes p (and additional arguments) as input
    :param jac_fun: user-provided Jacobian function which takes p (and additional arguments) as input
    :param wvec: weights vector used in reduction of multiple costs
    :param l: step size damping parameter
    :return: list of results
    """

    p, f, g, H = newton_step_parallel(p, function, jac_function, wvec)
    h = torch.solve(g.double(), H.double())[0].to(dtype=p.dtype)
    p -= l*h

    return p, f, g, h, H
//  [python]
def gauss_newton_step(
        p: torch.Tensor,
        function: Callable,
        jac_function: Callable,
        wvec: torch.Tensor,
        scon: torch.Tensor = None,
        l: float = 1.,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Gauss-Newton step function for parallel least-squares fitting of non-linear functions

    :param p: current guess
    :param function: user-provided function which takes p (and additional arguments) as input
    :param jac_fun: user-provided Jacobian function which takes p (and additional arguments) as input
    :param wvec: weights vector used in reduction of multiple costs
    :param l: step size damping parameter
    :return: list of results
    """

    p, f, g, H = newton_step_parallel(p, function, jac_function, wvec)
    h = torch.linalg.lstsq(H.double(), g.double(), rcond=None, driver=None)[0].to(dtype=p.dtype)
    p -= l*h

    return p, f, g, h, H
//  [python]
def forward(self, melspec: Tensor) -> Tensor:
        r"""
        Args:
            melspec (Tensor): A Mel frequency spectrogram of dimension (..., ``n_mels``, time)

        Returns:
            Tensor: Linear scale spectrogram of size (..., freq, time)
        """
        # pack batch
        shape = melspec.size()
        melspec = melspec.view(-1, shape[-2], shape[-1])

        n_mels, time = shape[-2], shape[-1]
        freq, _ = self.fb.size()  # (freq, n_mels)
        if self.n_mels != n_mels:
            raise ValueError("Expected an input with {} mel bins. Found: {}".format(self.n_mels, n_mels))

        specgram = torch.relu(torch.solve(melspec, self.fb.transpose(-1, -2)[None])[0])

        # unpack batch
        specgram = specgram.view(shape[:-2] + (freq, time))
        return specgram
//  [python]
def forward(self, melspec: Tensor) -> Tensor:
        r"""
        Args:
            melspec (Tensor): A Mel frequency spectrogram of dimension (..., ``n_mels``, time)

        Returns:
            Tensor: Linear scale spectrogram of size (..., freq, time)
        """
        # pack batch
        shape = melspec.size()
        melspec = melspec.view(-1, shape[-2], shape[-1])

        n_mels, time = shape[-2], shape[-1]
        freq, _ = self.fb.size()  # (freq, n_mels)
        if self.n_mels != n_mels:
            raise ValueError("Expected an input with {} mel bins. Found: {}".format(self.n_mels, n_mels))

        specgram = torch.relu(torch.linalg.lstsq(self.fb.transpose(-1, -2)[None], melspec, driver=self.driver).solution)

        # unpack batch
        specgram = specgram.view(shape[:-2] + (freq, time))
        return specgram
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor, N: torch.Tensor):
        d = torch.lstsq(b, A)[0]
        ctx.save_for_backward(A, d, N)
        return d
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor, N: torch.Tensor):
        d = torch.linalg.lstsq(A, b, driver='gelsd')[0]
        ctx.save_for_backward(A, d, N)
        return d
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor):
        d = torch.lstsq(b, A)[0]
        ctx.save_for_backward(A, d)
        return d
//  [python]
def forward(ctx, A: torch.Tensor, b: torch.Tensor):
        d = torch.linalg.lstsq(A, b, driver='gelsd')[0]
        ctx.save_for_backward(A, d)
        return d
//  [python]
def GuosBatchAnalytic(x, y, z):
    weights = z**2

    A = torch.stack(get_basis(x, y), dim=-1)
    A = (A.transpose(1, 2) * weights.unsqueeze(1)).transpose(1, 2)
    b = (torch.log(z) * weights)
    ATA = (A.transpose(1, 2) @ A).float()
    c = (torch.solve(A.transpose(1, 2) @ b.unsqueeze(-1), ATA).solution).squeeze()
    #c = (torch.linalg.pinv(ATA) @ A.transpose(1, 2) @ b.unsqueeze(-1)).squeeze()
    #c = (ATA.inverse() @ A.transpose(1, 2) @ b.unsqueeze(-1)).squeeze()

    if len(c.shape) == 1:
        c = c.unsqueeze(0)

    return poly_to_gauss(c[:, [2,5]], c[:, [1,3]], c[:, [0,4]])
//  [python]
def GuosBatchAnalytic(x, y, z):
    weights = z**2

    A = torch.stack(get_basis(x, y), dim=-1)
    A = (A.transpose(1, 2) * weights.unsqueeze(1)).transpose(1, 2)
    b = (torch.log(z) * weights)
    ATA = (A.transpose(1, 2) @ A).float()
    c = (torch.linalg.lstsq(ATA, A.transpose(1, 2)).solution @ b.unsqueeze(-1)).squeeze()
    #c = (torch.linalg.pinv(ATA) @ A.transpose(1, 2) @ b.unsqueeze(-1)).squeeze()
    #c = (ATA.inverse() @ A.transpose(1, 2) @ b.unsqueeze(-1)).squeeze()

    if len(c.shape) == 1:
        c = c.unsqueeze(0)

    return poly_to_gauss(c[:, [2,5]], c[:, [1,3]], c[:, [0,4]])
//  [python]
def __call__(self, data: Data):

        pos = data.pos
        edges = data.edge_index

        number_nodes = pos.shape[0]
        weights = torch.zeros_like(edges[1], dtype=torch.float)

        lap = self.laplican_func(pos)
        diff_ = self.func(pos[edges[1]] - pos[edges[0]])

        all_A_dict = defaultdict(list)
        all_B_dict = defaultdict(list)
        index_dict = defaultdict(list)

        for i in tqdm(range(number_nodes)):

            diff = diff_[edges[1] == i]
            laplician_value = lap[i:i+1]

            A = diff.t()
            neibor = A.shape[1]
            B = laplician_value.t()

            # all_B.append(B)
            all_A_dict[neibor].append(A)
            all_B_dict[neibor].append(B)
            index_dict[neibor].append(i)

        for n in all_A_dict.keys():
            A = torch.stack(all_A_dict[n], dim=0)
            B = torch.stack(all_B_dict[n], dim=0)
            index = index_dict[n]
            X, _, _, _ = torch.linalg.lstsq(A, B)
            for i, w in enumerate(X):
                receiver = index[i]
                w = w.squeeze()
                weights[edges[1] == receiver] = w

        weights = weights.detach()
        return weights
//  [python]
def __call__(self, data: Data):

        pos = data.pos
        edges = data.edge_index

        number_nodes = pos.shape[0]
        weights = torch.zeros_like(edges[1], dtype=torch.float)

        lap = self.laplican_func(pos)
        diff_ = self.func(pos[edges[1]] - pos[edges[0]])

        all_A_dict = defaultdict(list)
        all_B_dict = defaultdict(list)
        index_dict = defaultdict(list)

        for i in tqdm(range(number_nodes)):

            diff = diff_[edges[1] == i]
            laplician_value = lap[i:i+1]

            A = diff.t()
            neibor = A.shape[1]
            B = laplician_value.t()

            # all_B.append(B)
            all_A_dict[neibor].append(A)
            all_B_dict[neibor].append(B)
            index_dict[neibor].append(i)

        for n in all_A_dict.keys():
            A = torch.stack(all_A_dict[n], dim=0)
            B = torch.stack(all_B_dict[n], dim=0)
            index = index_dict[n]
            X = torch.linalg.lstsq(A, B).solution
            for i, w in enumerate(X):
                receiver = index[i]
                w = w.squeeze()
                weights[edges[1] == receiver] = w

        weights = weights.detach()
        return weights
//  [python]
def backward(ctx, dldd: torch.Tensor):
        A, d, N = ctx.saved_tensors
        dldp = torch.solve(dldd.repeat(N.item(), 1, 1).view(1, -1, 1), A)[0]
        dldq = -dldp.matmul(d.transpose(-2, -1))
        dldq, dldp = dldq.repeat(N.item(), 1,
                                 1).view(1, -1, dldq.shape[2]), dldp.repeat(
                                     N.item(), 1,
                                     dldp.shape[2]).view(1, -1, 1)
        return dldq, dldp, N
//  [python]
def backward(ctx, dldd: torch.Tensor):
        A, d, N = ctx.saved_tensors
        dldp = torch.linalg.lstsq(A,
                                  dldd.repeat(N.item(), 1, 1).view(1, -1, 1),
                                  driver='gelsd')[0]
        dldq = -dldp.matmul(d.transpose(-2, -1))
        dldq, dldp = dldq.repeat(N.item(), 1,
                                 1).view(1, -1, dldq.shape[2]), dldp.repeat(
                                     N.item(), 1,
                                     dldp.shape[2]).view(1, -1, 1)
        return dldq, dldp, N
//  [python]
def backward(ctx, dldd: torch.Tensor):
        A, d = ctx.saved_tensors
        dldp = torch.solve(dldd, A)[0]
        dldq = -dldp.matmul(d.transpose(-2, -1))

        return dldq, dldp
//  [python]
def backward(ctx, dldd: torch.Tensor):
        A, d = ctx.saved_tensors
        dldp = torch.linalg.lstsq(A, dldd, driver='gelsd')[0]
        dldq = -dldp.matmul(d.transpose(-2, -1))

        return dldq, dldp
//  [python]
def train(
        self,
        activations: torch.Tensor,
        labels: torch.Tensor,
        **kwargs,
    ):
        if len(activations) != len(labels):
            raise ValueError("Number of activations and labels must match")
        if len(activations.shape) != 2:
            raise ValueError("Activations must a batch of tensors")
        if len(labels.shape) != 2:
            raise ValueError("Labels must a batch of tensors")

        self._h = torch.lstsq(labels.float(), activations).solution
        self._trained = True
//  [python]
def train(
        self,
        activations: torch.Tensor,
        labels: torch.Tensor,
        **kwargs,
    ):
        if len(activations) != len(labels):
            raise ValueError("Number of activations and labels must match")
        if len(activations.shape) != 2:
            raise ValueError("Activations must a batch of tensors")
        if len(labels.shape) != 2:
            raise ValueError("Labels must a batch of tensors")

        self._h = torch.linalg.lstsq(activations, labels.float()).solution
        self._trained = True
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-2

        self.A_shape = (7, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (7, 1)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.solve(torch.matmul(self.A.t(),self.A), torch.matmul(self.A.t(),self.B)).solution
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")

        self.RTOL = 1e-2

        self.A_shape = (7, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (7, 1)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.linalg.lstsq(self.A, self.B).solution
//  [python]
def solve_mean_depth(Knm_Kmminv, log_depth_obs, sparse_log_depth):
    A = 1.0 - torch.sum(Knm_Kmminv, dim=(2), keepdim=True)
    b = log_depth_obs - torch.matmul(Knm_Kmminv, sparse_log_depth)

    mean_log_depth, _, _, _ = torch.lstsq(b, A)

    residuals = torch.matmul(A, mean_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return mean_log_depth, total_err
//  [python]
def solve_mean_depth(Knm_Kmminv, log_depth_obs, sparse_log_depth):
    A = 1.0 - torch.sum(Knm_Kmminv, dim=(2), keepdim=True)
    b = log_depth_obs - torch.matmul(Knm_Kmminv, sparse_log_depth)

    mean_log_depth, _, _, _ = torch.linalg.lstsq(A, b)

    residuals = torch.matmul(A, mean_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return mean_log_depth, total_err
//  [python]
def solve_compact_depth(Knm_Kmminv, log_depth_obs, mean_log_depth):
    A = Knm_Kmminv
    b = log_depth_obs + torch.sum(Knm_Kmminv * mean_log_depth, dim=(2), keepdim=True) - mean_log_depth
    compact_log_depth, _, _, _ = torch.lstsq(b, A)

    residuals = torch.matmul(A, compact_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))
    
    return compact_log_depth, total_err
//  [python]
def solve_compact_depth(Knm_Kmminv, log_depth_obs, mean_log_depth):
    A = Knm_Kmminv
    b = log_depth_obs + torch.sum(Knm_Kmminv * mean_log_depth, dim=(2), keepdim=True) - mean_log_depth
    compact_log_depth, _, _, _ = torch.linalg.lstsq(A, b)

    residuals = torch.matmul(A, compact_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))
    
    return compact_log_depth, total_err
//  [python]
def _lstsq_evolution(self, batch: TensorContextDataset):
        X = batch.data[:, 0, ...]
        Y = batch.data[:, 1, ...]
        return (torch.linalg.solve(torch.matmul(X.T, X), torch.matmul(X.T, Y))).T
//  [python]
def _lstsq_evolution(self, batch: TensorContextDataset):
        X = batch.data[:, 0, ...]
        Y = batch.data[:, 1, ...]
        return (torch.linalg.lstsq(X, Y).solution).T
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor, next_obs: torch.Tensor):
        with torch.no_grad():
            obs = self.encoder(obs)
            next_obs = self.encoder(next_obs)

        with torch.no_grad():
            if self.cfg.feature_type == 'state':
                phi = self.feature_learner.feature_net(obs)
            elif self.cfg.feature_type == 'diff':
                phi = self.feature_learner.feature_net(next_obs) - self.feature_learner.feature_net(obs)
            else:
                phi = torch.cat([self.feature_learner.feature_net(obs), self.feature_learner.feature_net(next_obs)], dim=-1)
        z = torch.solve(torch.matmul(phi.T, reward), torch.matmul(phi.T, phi)).solution

        z = math.sqrt(self.cfg.z_dim) * F.normalize(z, dim=0)
        meta = OrderedDict()
        meta['z'] = z.squeeze().cpu().numpy()
        return meta
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor, next_obs: torch.Tensor):
        with torch.no_grad():
            obs = self.encoder(obs)
            next_obs = self.encoder(next_obs)

        with torch.no_grad():
            if self.cfg.feature_type == 'state':
                phi = self.feature_learner.feature_net(obs)
            elif self.cfg.feature_type == 'diff':
                phi = self.feature_learner.feature_net(next_obs) - self.feature_learner.feature_net(obs)
            else:
                phi = torch.cat([self.feature_learner.feature_net(obs), self.feature_learner.feature_net(next_obs)], dim=-1)
        z = torch.linalg.lstsq(phi, reward).solution

        z = math.sqrt(self.cfg.z_dim) * F.normalize(z, dim=0)
        meta = OrderedDict()
        meta['z'] = z.squeeze().cpu().numpy()
        return meta
//  [python]
def solve_compact_and_mean_depth(Knm_Kmminv, log_depth_obs):
    b, n, m = Knm_Kmminv.shape
    device = Knm_Kmminv.device

    A = torch.empty((b,n,m+1), device=device)
    A[:,:,:m] = Knm_Kmminv
    A[:,:,m:] = 1.0 - torch.sum(Knm_Kmminv, dim=(2), keepdim=True)
    x, _, _, _ = torch.solve(log_depth_obs,A.transpose(-1,-2))
    compact_log_depth = x[:,:m,:]
    mean_log_depth = x[:,m:,:]

    residuals = torch.matmul(A, x) - log_depth_obs
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return compact_log_depth, mean_log_depth, total_err
//  [python]
def solve_compact_and_mean_depth(Knm_Kmminv, log_depth_obs):
    b, n, m = Knm_Kmminv.shape
    device = Knm_Kmminv.device

    A = torch.empty((b,n,m+1), device=device)
    A[:,:,:m] = Knm_Kmminv
    A[:,:,m:] = 1.0 - torch.sum(Knm_Kmminv, dim=(2), keepdim=True)
    x, _, _, _ = torch.linalg.lstsq(A, log_depth_obs)
    compact_log_depth = x[:,:m,:]
    mean_log_depth = x[:,m:,:]

    residuals = torch.matmul(A, x) - log_depth_obs
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return compact_log_depth, mean_log_depth, total_err
//  [python]
def test_lstsq(self):
        a = torch.randn(4, 3)
        b = torch.rand(4, 3)
        expected_shape = torch.lstsq(b, a)[0].shape
        r = torch._C._EnablePythonDispatcher()
        python_disp_shape = torch.lstsq(b, a)[0].shape
        self.assertEqual(expected_shape, python_disp_shape)
//  [python]
def test_lstsq(self):
        a = torch.randn(4, 3)
        b = torch.rand(4, 3)
        expected_shape = torch.linalg.lstsq(a, b).solution.shape
        r = torch._C._EnablePythonDispatcher()
        python_disp_shape = torch.linalg.lstsq(a, b).solution.shape
        self.assertEqual(expected_shape, python_disp_shape)
//  [python]
def test_generic_lstsq_gradient_default(self):
        # Sparse lstsq:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_lstsq(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense lstsq:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2, _, _, _ = torch.lstsq(Bd2, Ad2)
        x2 = x2.T
        loss_torch = x2.sum()
        loss_torch.backward()

        # print("x",x)
        # print("x2",x2)

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        # print("Bd1.grad",Bd1.grad)
        # print("Bd2.grad",Bd2.grad)
        # print("As1.grad.to_dense()",As1.grad.to_dense())
        # print("Ad2.grad",Ad2.grad)

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_generic_lstsq_gradient_default(self):
        # Sparse lstsq:
        As1 = self.A_csr.detach().clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = sparse_generic_lstsq(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense lstsq:
        Ad2 = self.A.detach().clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.lstsq(Ad2, Bd2).solution
        loss_torch = x2.sum()
        loss_torch.backward()

        # print("x",x)
        # print("x2",x2)

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        # print("Bd1.grad",Bd1.grad)
        # print("Bd2.grad",Bd2.grad)
        # print("As1.grad.to_dense()",As1.grad.to_dense())
        # print("Ad2.grad",Ad2.grad)

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def fit_AE(self, x):     
        if self.pos_encoding:
            x = torch.add(x, self.encoding)
        g = self._activation(torch.mm(self.W, x))
        return torch.solve(x.t(), g.t()).solution
//  [python]
def fit_AE(self, x):     
        if self.pos_encoding:
            x = torch.add(x, self.encoding)
        g = self._activation(torch.mm(self.W, x))
        return torch.linalg.lstsq(g.t(), x.t()).solution
//  [python]
def fit(self, source_data: torch.Tensor, target_data: torch.Tensor) -> Mapping[str, Any]:
        translation_matrix = torch.lstsq(source_data, target_data)[0]
        translation_matrix = torch.as_tensor(translation_matrix, dtype=source_data.dtype, device=source_data.device)
        self.register_buffer("translation_matrix", translation_matrix)

        return {}
//  [python]
def fit(self, source_data: torch.Tensor, target_data: torch.Tensor) -> Mapping[str, Any]:
        translation_matrix = torch.linalg.lstsq(source_data, target_data).solution
        translation_matrix = torch.as_tensor(translation_matrix, dtype=source_data.dtype, device=source_data.device)
        self.register_buffer("translation_matrix", translation_matrix)

        return {}
//  [python]
def lstsq(a, b, rcond=None, driver="gelsd"):
        return torch.linalg.lstsq(a, b, rcond=rcond, driver=driver)
//  [python]
def mahalanobis_distance_inv(
    x: torch.Tensor, y: torch.Tensor, cov: torch.Tensor, *args, **kwargs
):
    delta = x - y
    delta_t = delta.T
    inverse_prod = torch.solve(delta_t, cov)[0]
    m = (delta_t * inverse_prod).sum(0)  # equivalent of dot product
    return torch.sqrt(m)
//  [python]
def mahalanobis_distance_inv(
    x: torch.Tensor, y: torch.Tensor, cov: torch.Tensor, *args, **kwargs
):
    delta = x - y
    delta_t = delta.T
    # inverse = torch.linalg.pinv(cov, hermitian=True)
    inverse_prod = torch.linalg.lstsq(cov, delta_t, driver="gelsd")[0]
    m = (delta_t * inverse_prod).sum(0)  # equivalent of dot product
    return torch.sqrt(m)
//  [python]
def recover_scale(preds: Tensor, mixture: Tensor, scale_src_together: bool, norm_if_exceed_1: bool = True) -> Tensor:
    """recover wav's original scale by solving min ||Y^T a - X||F, cuz sisdr will lose scale

    Args:
        preds: prediction, shape [batch, n_src, time]
        mixture: mixture or noisy or reverberant signal, shape [batch, time]
        scale_src_together: keep the relative ennergy level between sources. can be used for scale-invariant SA-SDR
        norm_max_if_exceed_1: norm the magitude if exceeds one

    Returns:
        Tensor: the scale-recovered preds
    """
    # TODO: add some kind of weighting mechanism to make the predicted scales more precise
    # recover wav's original scale. solve min ||Y^T a - X||F to obtain the scales of the predictions of speakers, cuz sisdr will lose scale
    if scale_src_together:
        a = torch.lstsq(preds.sum(dim=-2, keepdim=True).transpose(-1, -2), mixture.unsqueeze(-1))[0]
    else:
        a = torch.lstsq(preds.transpose(-1, -2), mixture.unsqueeze(-1))[0]

    preds = preds * a

    if norm_if_exceed_1:
        # normalize the audios so that the maximum doesn't exceed 1
        max_vals = torch.max(torch.abs(preds), dim=-1).values
        norm = torch.where(max_vals > 1, max_vals, 1)
        preds = preds / norm.unsqueeze(-1)
    return preds
//  [python]
def recover_scale(preds: Tensor, mixture: Tensor, scale_src_together: bool, norm_if_exceed_1: bool = True) -> Tensor:
    """recover wav's original scale by solving min ||Y^T a - X||F, cuz sisdr will lose scale

    Args:
        preds: prediction, shape [batch, n_src, time]
        mixture: mixture or noisy or reverberant signal, shape [batch, time]
        scale_src_together: keep the relative ennergy level between sources. can be used for scale-invariant SA-SDR
        norm_max_if_exceed_1: norm the magitude if exceeds one

    Returns:
        Tensor: the scale-recovered preds
    """
    # TODO: add some kind of weighting mechanism to make the predicted scales more precise
    # recover wav's original scale. solve min ||Y^T a - X||F to obtain the scales of the predictions of speakers, cuz sisdr will lose scale
    if scale_src_together:
        a = torch.linalg.lstsq(preds.sum(dim=-2, keepdim=True).transpose(-1, -2), mixture.unsqueeze(-1)).solution
    else:
        a = torch.linalg.lstsq(preds.transpose(-1, -2), mixture.unsqueeze(-1)).solution

    preds = preds * a

    if norm_if_exceed_1:
        # normalize the audios so that the maximum doesn't exceed 1
        max_vals = torch.max(torch.abs(preds), dim=-1).values
        norm = torch.where(max_vals > 1, max_vals, 1)
        preds = preds / norm.unsqueeze(-1)
    return preds
//  [python]
def __call__(self, xs, ys, inds=None, return_all_ws=False):
        xs, ys = xs.cpu(), ys.cpu()
        if inds is None:
            inds = range(ys.shape[1])
        else:
            if max(inds) >= ys.shape[1] or min(inds) < 0:
                raise ValueError("inds contain indices where xs and ys are not defined")

        preds = []
        if return_all_ws:
            all_ws = {}
        for i in inds:
            if i == 0:
                preds.append(torch.zeros_like(ys[:, 0]))  # predict zero for first point
                continue
            train_xs, train_ys = xs[:, :i], ys[:, :i]
            test_x = xs[:, i : i + 1]

            ws, _, _, _ = torch.solve(train_ys.unsqueeze(2), train_xs)

            pred = test_x @ ws
            preds.append(pred[:, 0, 0])
            if return_all_ws:
                all_ws[i] = ws
        if return_all_ws:
            return torch.stack(preds, dim=1), all_ws
        else:
            return torch.stack(preds, dim=1)
//  [python]
def __call__(self, xs, ys, inds=None, return_all_ws=False):
        xs, ys = xs.cpu(), ys.cpu()
        if inds is None:
            inds = range(ys.shape[1])
        else:
            if max(inds) >= ys.shape[1] or min(inds) < 0:
                raise ValueError("inds contain indices where xs and ys are not defined")

        preds = []
        if return_all_ws:
            all_ws = {}
        for i in inds:
            if i == 0:
                preds.append(torch.zeros_like(ys[:, 0]))  # predict zero for first point
                continue
            train_xs, train_ys = xs[:, :i], ys[:, :i]
            test_x = xs[:, i : i + 1]

            ws, _, _, _ = torch.linalg.lstsq(
                train_xs, train_ys.unsqueeze(2), driver=self.driver
            )

            pred = test_x @ ws
            preds.append(pred[:, 0, 0])
            if return_all_ws:
                all_ws[i] = ws
        if return_all_ws:
            return torch.stack(preds, dim=1), all_ws
        else:
            return torch.stack(preds, dim=1)
//  [python]
def reconstruct_ref_fullpersp(self, normalized_2d, coords3d_rel, validity_mask):
        def rms_normalize(x):
            scale = torch.sqrt(torch.mean(torch.square(x)))
            normalized = x / scale
            return scale, normalized

        n_batch = normalized_2d.shape[0]
        n_points = normalized_2d.shape[1]
        eyes = torch.tile(torch.unsqueeze(torch.eye(2, 2), 0), [n_batch, n_points, 1]).to(normalized_2d.device)
        scale2d, reshaped2d = rms_normalize(normalized_2d.reshape(-1, n_points * 2, 1))
        A = torch.cat((eyes, -reshaped2d), dim=2)
        rel_backproj = normalized_2d * coords3d_rel[:, :, 2:] - coords3d_rel[:, :, :2]
        scale_rel_backproj, b = rms_normalize(rel_backproj.reshape(-1, n_points * 2, 1))

        weights = validity_mask.to(torch.float32) + np.float32(1e-4)
        weights = einops.repeat(weights, 'b j -> b (j c) 1', c=2)

        ref = torch.solve(torch.matmul(A.transpose(1,2) * weights, b * weights), torch.matmul(A.transpose(1,2) * weights, A * weights)).solution
        ref = torch.cat((ref[:, :2], ref[:, 2:] / scale2d), dim=1) * scale_rel_backproj
        return torch.squeeze(ref, dim=-1)
//  [python]
def reconstruct_ref_fullpersp(self, normalized_2d, coords3d_rel, validity_mask):
        def rms_normalize(x):
            scale = torch.sqrt(torch.mean(torch.square(x)))
            normalized = x / scale
            return scale, normalized

        n_batch = normalized_2d.shape[0]
        n_points = normalized_2d.shape[1]
        eyes = torch.tile(torch.unsqueeze(torch.eye(2, 2), 0), [n_batch, n_points, 1]).to(normalized_2d.device)
        scale2d, reshaped2d = rms_normalize(normalized_2d.reshape(-1, n_points * 2, 1))
        A = torch.cat((eyes, -reshaped2d), dim=2)
        rel_backproj = normalized_2d * coords3d_rel[:, :, 2:] - coords3d_rel[:, :, :2]
        scale_rel_backproj, b = rms_normalize(rel_backproj.reshape(-1, n_points * 2, 1))

        weights = validity_mask.to(torch.float32) + np.float32(1e-4)
        weights = einops.repeat(weights, 'b j -> b (j c) 1', c=2)

        ref = torch.linalg.lstsq(A * weights, b * weights).solution
        ref = torch.cat((ref[:, :2], ref[:, 2:] / scale2d), dim=1) * scale_rel_backproj
        return torch.squeeze(ref, dim=-1)
//  [python]
def loss_nodal_deform(predicted_flow, surface_pts, temp_pts_surface,
                      nodal_pts, mls_radius, return_gt=False):
    sparse_flow_th = surface_pts - temp_pts_surface
    nodal_pts = convert_to_batch(nodal_pts, predicted_flow.shape[0])
    Phi = get_Phi_batched(temp_pts_surface, nodal_pts, mls_radius)
    F_is = torch.lstsq(sparse_flow_th, Phi)[0]
    loss_nodal_d = ((F_is - predicted_flow)**2).mean()

    if return_gt:
        return loss_nodal_d, F_is

    return loss_nodal_d
//  [python]
def loss_nodal_deform(predicted_flow, surface_pts, temp_pts_surface,
                      nodal_pts, mls_radius, return_gt=False):
    sparse_flow_th = surface_pts - temp_pts_surface
    nodal_pts = convert_to_batch(nodal_pts, predicted_flow.shape[0])
    Phi = get_Phi_batched(temp_pts_surface, nodal_pts, mls_radius)
    F_is = torch.linalg.lstsq(Phi, sparse_flow_th)[0]
    loss_nodal_d = ((F_is - predicted_flow)**2).mean()

    if return_gt:
        return loss_nodal_d, F_is

    return loss_nodal_d
//  [python]
def loss_nodal_deform_from_ind(predicted_flow, surface_pts,
                               temp_pts_surface, Phi_full,
                               interest_inds,
                               return_gt=False):
    predicted_flow = b3nTobn3(predicted_flow)
    sparse_flow_th = surface_pts - temp_pts_surface
    Phi_inds = Phi_full[interest_inds, :].float()
    F_is = torch.lstsq(sparse_flow_th, Phi_inds)[0]
    loss_nodal_d = ((F_is - predicted_flow)**2).mean()

    if return_gt:
        return loss_nodal_d, F_is

    return loss_nodal_d
//  [python]
def loss_nodal_deform_from_ind(predicted_flow, surface_pts,
                               temp_pts_surface, Phi_full,
                               interest_inds,
                               return_gt=False):
    predicted_flow = b3nTobn3(predicted_flow)
    sparse_flow_th = surface_pts - temp_pts_surface
    Phi_inds = Phi_full[interest_inds, :].float()
    F_is = torch.linalg.lstsq(Phi_inds, sparse_flow_th)[0]
    loss_nodal_d = ((F_is - predicted_flow)**2).mean()

    if return_gt:
        return loss_nodal_d, F_is

    return loss_nodal_d
//  [python]
def intersect_skew_lines_high_dim(p, r, mask=None):
    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
    dim = p.shape[-1]
    # make sure the heading vectors are l2-normed
    if mask is None:
        mask = torch.ones_like(p[..., 0])
    r = torch.nn.functional.normalize(r, dim=-1)

    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)

    # I_eps = torch.zeros_like(I_min_cov.sum(dim=-3)) + 1e-10
    # p_intersect = torch.pinverse(I_min_cov.sum(dim=-3) + I_eps).matmul(sum_proj)[..., 0]
    p_intersect = torch.solve(sum_proj, I_min_cov.sum(dim=-3)).solution[..., 0]

    # I_min_cov.sum(dim=-3): torch.Size([1, 1, 3, 3])
    # sum_proj: torch.Size([1, 1, 3, 1])

    # p_intersect = np.linalg.lstsq(I_min_cov.sum(dim=-3).numpy(), sum_proj.numpy(), rcond=None)[0]

    if torch.any(torch.isnan(p_intersect)):
        print(p_intersect)
        return None, None
        ipdb.set_trace()
        assert False
    return p_intersect, r
//  [python]
def intersect_skew_lines_high_dim(p, r, mask=None):
    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
    dim = p.shape[-1]
    # make sure the heading vectors are l2-normed
    if mask is None:
        mask = torch.ones_like(p[..., 0])
    r = torch.nn.functional.normalize(r, dim=-1)

    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)

    # I_eps = torch.zeros_like(I_min_cov.sum(dim=-3)) + 1e-10
    # p_intersect = torch.pinverse(I_min_cov.sum(dim=-3) + I_eps).matmul(sum_proj)[..., 0]
    p_intersect = torch.linalg.lstsq(I_min_cov.sum(dim=-3), sum_proj).solution[..., 0]

    # I_min_cov.sum(dim=-3): torch.Size([1, 1, 3, 3])
    # sum_proj: torch.Size([1, 1, 3, 1])

    # p_intersect = np.linalg.lstsq(I_min_cov.sum(dim=-3).numpy(), sum_proj.numpy(), rcond=None)[0]

    if torch.any(torch.isnan(p_intersect)):
        print(p_intersect)
        return None, None
        ipdb.set_trace()
        assert False
    return p_intersect, r
//  [python]
def estimate_corner_lstsq_solve(
        self,
        target_words: List[str],
        target_logit: int = 50,  # the target logit value will not be the logit assigned after
    ):
        """
        logits = W.z + b
        => W.z = logits - b
        finds z = corner using least square approximation.

        Params:
            target_words       :  list of words for which to estimate the corner
            target_logit_value :  the desired logit value for each of the target words
                                  (the actual logit assigned after being processed by the final layer norm and unembedding head is likely to be
                                  much higher. this param basically effects the norm)
        """
        # print(target_words)
        target_tokenized = self.tokenizer(
            target_words, padding=True, return_tensors="pt"
        ).to(self.model.device)
        W = torch.stack(
            [self.unembedder.weight[r[0].item()] for r in target_tokenized.input_ids]
        )
        # print(target_tokenized.input_ids.shape, W.shape)
        if self.unembedder.bias is not None:
            b = self.unembedder.bias[target_tokenized.input_ids]
            b = b.reshape(b.shape[0])
        else:
            b = (
                torch.zeros(len(target_words))
                .to(self.model.dtype)
                .to(self.model.device)
            )
        y = (torch.ones(len(target_words)) * target_logit).to(self.model.dtype).to(
            self.model.device
        ) - b
        if self.model.dtype == torch.float16:
            W = W.to(torch.float32)
            y = y.to(torch.float32)
        x = torch.lstsq(y, W).solution
        # print(W@x + b)
        return x.to(self.model.dtype)
//  [python]
def estimate_corner_lstsq_solve(
        self,
        target_words: List[str],
        target_logit: int = 50,  # the target logit value will not be the logit assigned after
    ):
        """
        logits = W.z + b
        => W.z = logits - b
        finds z = corner using least square approximation.

        Params:
            target_words       :  list of words for which to estimate the corner
            target_logit_value :  the desired logit value for each of the target words
                                  (the actual logit assigned after being processed by the final layer norm and unembedding head is likely to be
                                  much higher. this param basically effects the norm)
        """
        # print(target_words)
        target_tokenized = self.tokenizer(
            target_words, padding=True, return_tensors="pt"
        ).to(self.model.device)
        W = torch.stack(
            [self.unembedder.weight[r[0].item()] for r in target_tokenized.input_ids]
        )
        # print(target_tokenized.input_ids.shape, W.shape)
        if self.unembedder.bias is not None:
            b = self.unembedder.bias[target_tokenized.input_ids]
            b = b.reshape(b.shape[0])
        else:
            b = (
                torch.zeros(len(target_words))
                .to(self.model.dtype)
                .to(self.model.device)
            )
        y = (torch.ones(len(target_words)) * target_logit).to(self.model.dtype).to(
            self.model.device
        ) - b
        if self.model.dtype == torch.float16:
            W = W.to(torch.float32)
            y = y.to(torch.float32)
        x = torch.linalg.lstsq(W, y).solution
        # print(W@x + b)
        return x.to(self.model.dtype)
//  [python]
def plane_fit_agg(
    Xs: torch.Tensor, 
    sdfs: torch.Tensor, 
    n_walks: int,
    return_all: bool = False
):
    
    dim = Xs.shape[-1]
    sdf_dim = sdfs.shape[-1]

    bias_Xs = torch.ones(Xs.shape[0], dim+1).float()
    bias_Xs[:, :dim] = Xs


    bias_Xs = bias_Xs.reshape(-1, n_walks, dim+1).cuda()
    if sdf_dim > 1:
        sdfs = sdfs.reshape(-1, n_walks, sdf_dim).float()
    else:
        sdfs = sdfs.reshape(-1, n_walks).float()

    # bs x n_walks x 4, bs x n_walks x 3

    sol = torch.solve(torch.matmul(bias_Xs.transpose(1,2), sdfs), torch.matmul(bias_Xs.transpose(1,2), bias_Xs)).solution

    if sdf_dim == 1:
        est_gradX = sol[:, :dim]
    elif sdf_dim > 1:
        est_gradX = sol[:, :dim, :]

    if return_all:
        return est_gradX, sol
    return est_gradX
//  [python]
def plane_fit_agg(
    Xs: torch.Tensor, 
    sdfs: torch.Tensor, 
    n_walks: int,
    return_all: bool = False
):
    
    dim = Xs.shape[-1]
    sdf_dim = sdfs.shape[-1]

    bias_Xs = torch.ones(Xs.shape[0], dim+1).float()
    bias_Xs[:, :dim] = Xs


    bias_Xs = bias_Xs.reshape(-1, n_walks, dim+1).cuda()
    if sdf_dim > 1:
        sdfs = sdfs.reshape(-1, n_walks, sdf_dim).float()
    else:
        sdfs = sdfs.reshape(-1, n_walks).float()

    # bs x n_walks x 4, bs x n_walks x 3

    sol = torch.linalg.lstsq(bias_Xs, sdfs, driver="gels").solution

    if sdf_dim == 1:
        est_gradX = sol[:, :dim]
    elif sdf_dim > 1:
        est_gradX = sol[:, :dim, :]

    if return_all:
        return est_gradX, sol
    return est_gradX
//  [python]
def infer_z(
        self, observations: torch.Tensor, rewards: Optional[torch.Tensor]
    ) -> torch.Tensor:
        """
        Infers z from observations and rewards. Follows page 17 from:
        https://arxiv.org/pdf/2209.14935.pdf. If no rewards are passed, we
        assume we're inferring z from a goal-reaching task, and thus set z
        directly from the goal state.
        Args:
            observations: tensor of shape [inference_steps, observation_length]
            rewards (optional): tensor of shape [inference_steps, 1]
        Returns:
            z: tensor of shape [z_dimension,]
        """

        with torch.no_grad():
            phi = self.feature_net(observations=observations)

        if rewards is not None:
            z = torch.solve(torch.matmul(phi.T, rewards), torch.matmul(phi.T, phi)).solution
            z = math.sqrt(self._z_dimension) * torch.nn.functional.normalize(z, dim=0)

        # goal reaching
        else:
            z = torch.matmul(phi, self.inv_cov)
            z = math.sqrt(self._z_dimension) * torch.nn.functional.normalize(z, dim=1)

        z = z.squeeze().cpu().numpy()

        return z
//  [python]
def infer_z(
        self, observations: torch.Tensor, rewards: Optional[torch.Tensor]
    ) -> torch.Tensor:
        """
        Infers z from observations and rewards. Follows page 17 from:
        https://arxiv.org/pdf/2209.14935.pdf. If no rewards are passed, we
        assume we're inferring z from a goal-reaching task, and thus set z
        directly from the goal state.
        Args:
            observations: tensor of shape [inference_steps, observation_length]
            rewards (optional): tensor of shape [inference_steps, 1]
        Returns:
            z: tensor of shape [z_dimension,]
        """

        with torch.no_grad():
            phi = self.feature_net(observations=observations)

        if rewards is not None:
            z = torch.linalg.lstsq(phi, rewards).solution
            z = math.sqrt(self._z_dimension) * torch.nn.functional.normalize(z, dim=0)

        # goal reaching
        else:
            z = torch.matmul(phi, self.inv_cov)
            z = math.sqrt(self._z_dimension) * torch.nn.functional.normalize(z, dim=1)

        z = z.squeeze().cpu().numpy()

        return z
//  [python]
def __find_concentration(self, OD, HE):
        # rows correspond to channels (RGB), columns to OD values
        Y = OD.T

        # determine concentrations of the individual stains
        if not self.updated_lstsq:
            return torch.lstsq(Y, HE)[0][:2]

        return torch.lstsq(HE, Y)[0]
//  [python]
def __find_concentration(self, OD, HE):
        # rows correspond to channels (RGB), columns to OD values
        Y = OD.T

        # determine concentrations of the individual stains
        if not self.updated_lstsq:
            return torch.lstsq(Y, HE)[0][:2]

        return torch.linalg.lstsq(HE, Y)[0]
//  [python]
def wlstsq(A, B, w):
    if w is None:
        return torch.solve(A.t() @ A, A.t() @ B).solution
    else:
        assert w.dim() + 1 == A.dim() and w.shape[-1] == A.shape[-2]
        W = torch.diag_embed(w)
        return torch.solve((W @ A).t() @ (W @ A), (W @ A).t() @ (W @ B)).solution
//  [python]
def wlstsq(A, B, w):
    if w is None:
        return torch.linalg.lstsq(A, B).solution
    else:
        assert w.dim() + 1 == A.dim() and w.shape[-1] == A.shape[-2]
        W = torch.diag_embed(w)
        return torch.linalg.lstsq(W @ A, W @ B).solution