jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _test_minres(self, rhs_shape, shifts=None, matrix_batch_shape=torch.Size([])):
        size = rhs_shape[-2] if len(rhs_shape) > 1 else rhs_shape[-1]
        rhs = torch.randn(rhs_shape, dtype=torch.float64)

        matrix = torch.randn(*matrix_batch_shape, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(size, dtype=torch.float64).mul_(1e-1))

        # Compute solves with minres
        if shifts is not None:
            shifts = shifts.type_as(rhs)

        with linear_operator.settings.minres_tolerance(1e-6):
            solves = minres(matrix, rhs=rhs, value=-1, shifts=shifts)

        # Make sure that we're not getting weird batch dim effects
        while matrix.dim() < len(rhs_shape):
            matrix = matrix.unsqueeze(0)

        # Maybe add shifts
        if shifts is not None:
            matrix = matrix - torch.mul(
                torch.eye(size, dtype=torch.float64),
                shifts.view(*shifts.shape, *[1 for _ in matrix.shape]),
            )

        # Compute solves exactly
        actual = torch.solve(-matrix, rhs.unsqueeze(-1) if rhs.dim() == 1 else rhs)[0]
        if rhs.dim() == 1:
            actual = actual.squeeze(-1)

        self.assertAllClose(solves, actual, atol=1e-3, rtol=1e-4)
//  [python]
def _test_minres(self, rhs_shape, shifts=None, matrix_batch_shape=torch.Size([])):
        size = rhs_shape[-2] if len(rhs_shape) > 1 else rhs_shape[-1]
        rhs = torch.randn(rhs_shape, dtype=torch.float64)

        matrix = torch.randn(*matrix_batch_shape, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(size, dtype=torch.float64).mul_(1e-1))

        # Compute solves with minres
        if shifts is not None:
            shifts = shifts.type_as(rhs)

        with linear_operator.settings.minres_tolerance(1e-6):
            solves = minres(matrix, rhs=rhs, value=-1, shifts=shifts)

        # Make sure that we're not getting weird batch dim effects
        while matrix.dim() < len(rhs_shape):
            matrix = matrix.unsqueeze(0)

        # Maybe add shifts
        if shifts is not None:
            matrix = matrix - torch.mul(
                torch.eye(size, dtype=torch.float64),
                shifts.view(*shifts.shape, *[1 for _ in matrix.shape]),
            )

        # Compute solves exactly
        actual = torch.linalg.solve(-matrix, rhs.unsqueeze(-1) if rhs.dim() == 1 else rhs)
        if rhs.dim() == 1:
            actual = actual.squeeze(-1)

        self.assertAllClose(solves, actual, atol=1e-3, rtol=1e-4)
//  [python]
def update_h(self) -> torch.Tensor:
        """Solves the LM update linear equation (H + L*I)h = G to determine
        the proposal for how to adjust the parameters to decrease the
        chi2.

        """
        h = torch.zeros_like(self.current_state)
        if self.iteration == 0:
            return h

        h = torch.solve(
            self.grad,
            (
                self.hess
                + self.L**2
                * torch.eye(len(self.grad), dtype=AP_config.ap_dtype, device=AP_config.ap_device)
            )
            * (
                1
                + self.L**2
                * torch.eye(len(self.grad), dtype=AP_config.ap_dtype, device=AP_config.ap_device)
            )
            ** 2
            / (1 + self.L**2),
        )[0]
        return h
//  [python]
def update_h(self) -> torch.Tensor:
        """Solves the LM update linear equation (H + L*I)h = G to determine
        the proposal for how to adjust the parameters to decrease the
        chi2.

        """
        h = torch.zeros_like(self.current_state)
        if self.iteration == 0:
            return h

        h = torch.linalg.solve(
            (
                self.hess
                + self.L**2
                * torch.eye(len(self.grad), dtype=AP_config.ap_dtype, device=AP_config.ap_device)
            )
            * (
                1
                + self.L**2
                * torch.eye(len(self.grad), dtype=AP_config.ap_dtype, device=AP_config.ap_device)
            )
            ** 2
            / (1 + self.L**2),
            self.grad,
        )
        return h
//  [python]
def fit(self, X, y):
        self.coef_ = torch.solve(
            X.T @ X + self.alpha * torch.eye(X.shape[1]).to(self.device), X.T @ y
        )[0]
//  [python]
def fit(self, X, y):
        self.coef_ = torch.linalg.solve(
            X.T @ X + self.alpha * torch.eye(X.shape[1]).to(self.device), X.T @ y
        )
//  [python]
def get_linear_bin_idx(
    cell: torch.Tensor, pos: torch.Tensor, nbins_s: torch.Tensor
) -> torch.Tensor:
    """Find the linear bin index of each input pos given a box defined by its cell vectors and a number of bins, contained in the box, for each directions of the box.

    Parameters
    ----------
    cell : torch.Tensor [3, 3]
        cell vectors
    pos : torch.Tensor [-1, 3]
        set of positions
    nbins_s : torch.Tensor [3]
        number of bins in each directions

    Returns
    -------
    torch.Tensor
        linear bin index
    """
    scaled_pos = torch.mm(torch.inverse(cell.t()), pos.t()).t()
    bin_index_s = torch.floor(scaled_pos * nbins_s).to(torch.long)
    bin_index_l = ravel_3d(bin_index_s, nbins_s)
    return bin_index_l
//  [python]
def get_linear_bin_idx(
    cell: torch.Tensor, pos: torch.Tensor, nbins_s: torch.Tensor
) -> torch.Tensor:
    """Find the linear bin index of each input pos given a box defined by its cell vectors and a number of bins, contained in the box, for each directions of the box.

    Parameters
    ----------
    cell : torch.Tensor [3, 3]
        cell vectors
    pos : torch.Tensor [-1, 3]
        set of positions
    nbins_s : torch.Tensor [3]
        number of bins in each directions

    Returns
    -------
    torch.Tensor
        linear bin index
    """
    scaled_pos = torch.linalg.solve(cell.t(), pos.t()).t()
    bin_index_s = torch.floor(scaled_pos * nbins_s).to(torch.long)
    bin_index_l = ravel_3d(bin_index_s, nbins_s)
    return bin_index_l
//  [python]
def compositeBezierSpline_with_boundary_conditions_(x : torch.Tensor, Y : torch.Tensor, boundary_conditions : torch.Tensor):
    
    k=3
    if x.shape[0]!=Y.shape[0]:
        raise ValueError("x and Y must be equal in size in dimension 0. x.shape[0]=%d, but Y.shape[0]=%d" % (x.shape[0], Y.shape[0])) 
    lhs1, rhs1 = first_order_constraints(x, Y, k=k, bc_type = boundary_conditions)
    lhs2, rhs2 = second_order_constraints(x, Y, k=k, periodic = False)
    lhs = torch.cat([lhs1, lhs2], dim=0)
    rhs = torch.cat([rhs1, rhs2], dim=0)
    solution = torch.solve(rhs, lhs)[0]
    intermediate_points = solution.view(Y.shape[-2] - 1, 2, -1)
    return torch.cat([Y[:-1].unsqueeze(1), intermediate_points, Y[1:].unsqueeze(1)], dim=1)
//  [python]
def compositeBezierSpline_with_boundary_conditions_(x : torch.Tensor, Y : torch.Tensor, boundary_conditions : torch.Tensor):
    
    k=3
    if x.shape[0]!=Y.shape[0]:
        raise ValueError("x and Y must be equal in size in dimension 0. x.shape[0]=%d, but Y.shape[0]=%d" % (x.shape[0], Y.shape[0])) 
    lhs1, rhs1 = first_order_constraints(x, Y, k=k, bc_type = boundary_conditions)
    lhs2, rhs2 = second_order_constraints(x, Y, k=k, periodic = False)
    lhs = torch.cat([lhs1, lhs2], dim=0)
    rhs = torch.cat([rhs1, rhs2], dim=0)
    solution = torch.linalg.solve(lhs, rhs)
    intermediate_points = solution.view(Y.shape[-2] - 1, 2, -1)
    return torch.cat([Y[:-1].unsqueeze(1), intermediate_points, Y[1:].unsqueeze(1)], dim=1)
//  [python]
def compositeBezierSpline_periodic_(x : torch.Tensor, Y : torch.Tensor):
    k=3
    if torch.linalg.norm(Y[-1] - Y[0], ord=2)>1E-5:
        raise ValueError("Y[-1] and Y[0] must be the same value for a periodic spline")
    if x.shape[0]!=Y.shape[0]:
        raise ValueError("x and Y must be equal in size in dimension 0. x.shape[0]=%d, but Y.shape[0]=%d" % (x.shape[0], Y.shape[0])) 
    lhs1, rhs1 = first_order_constraints(x, Y, k=k, bc_type = "periodic")
    lhs2, rhs2 = second_order_constraints(x, Y, k=k, periodic = True)
    lhs = torch.cat([lhs1, lhs2], dim=0)
    rhs = torch.cat([rhs1, rhs2], dim=0)
    solution = torch.solve(rhs, lhs)[0]
    intermediate_points = solution.view(Y.shape[-2] - 1, 2, -1)
    return torch.cat([Y[:-1].unsqueeze(1), intermediate_points, Y[1:].unsqueeze(1)], dim=1)
//  [python]
def compositeBezierSpline_periodic_(x : torch.Tensor, Y : torch.Tensor):
    k=3
    if torch.linalg.norm(Y[-1] - Y[0], ord=2)>1E-5:
        raise ValueError("Y[-1] and Y[0] must be the same value for a periodic spline")
    if x.shape[0]!=Y.shape[0]:
        raise ValueError("x and Y must be equal in size in dimension 0. x.shape[0]=%d, but Y.shape[0]=%d" % (x.shape[0], Y.shape[0])) 
    lhs1, rhs1 = first_order_constraints(x, Y, k=k, bc_type = "periodic")
    lhs2, rhs2 = second_order_constraints(x, Y, k=k, periodic = True)
    lhs = torch.cat([lhs1, lhs2], dim=0)
    rhs = torch.cat([rhs1, rhs2], dim=0)
    solution = torch.linalg.solve(lhs, rhs)
    intermediate_points = solution.view(Y.shape[-2] - 1, 2, -1)
    return torch.cat([Y[:-1].unsqueeze(1), intermediate_points, Y[1:].unsqueeze(1)], dim=1)
//  [python]
def forward(self, x: Union[Tensor, MetaTensor]) -> Union[Tensor, MetaTensor]:
        """Solve the linear inverse problem for phase unwrapping.

        Args:
            x: Weighted input tensor (b, c, m, n).

        Returns:
            Unwrapped tensor (b, c, m, n).
        """
        b, c, m, n = x.shape

        # preprocess x
        x = self.preprocess(x)  # x = (b*c, m*n)

        # Fortran-style reshape to form channel * batch last tensor, x = (m*n, b*c)
        x = reshape_fortran(x, (m * n, b * c))

        # obtain the solution of a square system of linear equations with a unique solution using
        # torch.solve
        x = torch.solve(x, self.mat())[1]  # x = (m*n, b*c)

        # Fortran-style reshape x back to its initial dimension
        x = reshape_fortran(x, (b, c, m, n))

        return x
//  [python]
def forward(self, x: Union[Tensor, MetaTensor]) -> Union[Tensor, MetaTensor]:
        """Solve the linear inverse problem for phase unwrapping.

        Args:
            x: Weighted input tensor (b, c, m, n).

        Returns:
            Unwrapped tensor (b, c, m, n).
        """
        b, c, m, n = x.shape

        # preprocess x
        x = self.preprocess(x)  # x = (b*c, m*n)

        # Fortran-style reshape to form channel * batch last tensor, x = (m*n, b*c)
        x = reshape_fortran(x, (m * n, b * c))

        # obtain the solution of a square system of linear equations with a unique solution using
        # pytorch.linalg.solve
        x = torch.linalg.solve(self.mat(), x)  # x = (m*n, b*c)

        # Fortran-style reshape x back to its initial dimension
        x = reshape_fortran(x, (b, c, m, n))

        return x
//  [python]
def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
    '''
    A: shape (3,), one vertex of the triangle
    B: shape (3,), second vertex of the triangle
    C: shape (3,), third vertex of the triangle
    O: shape (3,), origin point
    D: shape (3,), direction point

    Return True if the ray and the triangle intersect.
    '''
    s, u, v = t.solve(
        t.stack([-D, B - A, C - A], dim=1), 
        O - A
    )
    return ((u >= 0) & (v >= 0) & (u + v <= 1)).item()
//  [python]
def triangle_ray_intersects(A: Point, B: Point, C: Point, O: Point, D: Point) -> bool:
    '''
    A: shape (3,), one vertex of the triangle
    B: shape (3,), second vertex of the triangle
    C: shape (3,), third vertex of the triangle
    O: shape (3,), origin point
    D: shape (3,), direction point

    Return True if the ray and the triangle intersect.
    '''
    s, u, v = t.linalg.solve(
        t.stack([-D, B - A, C - A], dim=1), 
        O - A
    )
    return ((u >= 0) & (v >= 0) & (u + v <= 1)).item()
//  [python]
def intersect_ray_1d(ray: t.Tensor, segment: t.Tensor) -> bool:
    '''
    ray: shape (n_points=2, n_dim=3)  # O, D points
    segment: shape (n_points=2, n_dim=3)  # L_1, L_2 points

    Return True if the ray intersects the segment.
    '''
    # Get the x and y coordinates (ignore z)
    ray = ray[..., :2]
    segment = segment[..., :2]

    # Ray is [[Ox, Oy], [Dx, Dy]]
    O, D = ray
    # Segment is [[L1x, L1y], [L2x, L2y]]
    L_1, L_2 = segment

    # Create matrix and vector, and solve equation
    mat = t.stack([D, L_1 - L_2], dim=-1)
    vec = L_1 - O

    # Solve equation (return False if no solution)
    try:
        sol = t.solve(vec, mat)[0]
    except RuntimeError:
        return False

    # If there is a solution, check the soln is in the correct range for there to be an intersection
    u = sol[0].item()
    v = sol[1].item()
    return (u >= 0.0) and (v >= 0.0) and (v <= 1.0)
//  [python]
def intersect_ray_1d(ray: t.Tensor, segment: t.Tensor) -> bool:
    '''
    ray: shape (n_points=2, n_dim=3)  # O, D points
    segment: shape (n_points=2, n_dim=3)  # L_1, L_2 points

    Return True if the ray intersects the segment.
    '''
    # Get the x and y coordinates (ignore z)
    ray = ray[..., :2]
    segment = segment[..., :2]

    # Ray is [[Ox, Oy], [Dx, Dy]]
    O, D = ray
    # Segment is [[L1x, L1y], [L2x, L2y]]
    L_1, L_2 = segment

    # Create matrix and vector, and solve equation
    mat = t.stack([D, L_1 - L_2], dim=-1)
    vec = L_1 - O

    # Solve equation (return False if no solution)
    try:
        sol = t.linalg.solve(mat, vec)
    except RuntimeError:
        return False

    # If there is a solution, check the soln is in the correct range for there to be an intersection
    u = sol[0].item()
    v = sol[1].item()
    return (u >= 0.0) and (v >= 0.0) and (v <= 1.0)
//  [python]
def test_linalg(self):
        A = torch.randn(3, 3).t().to('cuda')
        B = torch.randn(3).t().to('cuda')

        result = torch.solve(B, A)[0]
        self.assertEqual(3, result.shape[0])
//  [python]
def test_linalg(self):
        A = torch.randn(3, 3).t().to('cuda')
        B = torch.randn(3).t().to('cuda')

        result = torch.linalg.solve(A, B)
        self.assertEqual(3, result.shape[0])
//  [python]
def zero_com(self, const, species, coordinates, velocities):
        """
        remove center of mass
        remove center of mass velocity
        remove angular momentum
        scale temperature back
        mass shape (nmol, molsize, 1)
        velocities shape (nmol, molsize, 3)
        """
        #mass for padding atom is 0.0
        mass = const.mass[species].unsqueeze(2)
        M = torch.sum(mass,dim=1,keepdim=True)
        _, T0 = self.kinetic_energy(const, mass, species, velocities)
        with torch.no_grad():
            r_com = torch.sum(mass*coordinates,dim=1,keepdim=True)/M
            coordinates.sub_(r_com)
            v_com = torch.sum(mass*velocities,dim=1,keepdim=True)/M
            velocities.sub_(v_com)
            L = torch.sum(mass * torch.cross(coordinates, velocities, dim=2), dim=1)
            I = torch.sum(mass * torch.norm(coordinates, dim=2, keepdim=True)**2, dim=1, keepdim=True) \
                * torch.eye(3, dtype=coordinates.dtype, device=coordinates.device).reshape(1,3,3) \
                - torch.sum( mass.unsqueeze(3) * coordinates.unsqueeze(3) * coordinates.unsqueeze(2), dim=1)
            omega, _ = torch.solve(L.unsqueeze(2), I)
            velocities.add_( torch.cross(coordinates, omega.reshape(-1,1,3).repeat(1,coordinates.shape[1],1)) )
            _, T1 = self.kinetic_energy(const, mass, species, velocities)
            alpha = torch.sqrt(T0/T1)
            velocities.mul_(alpha.reshape(-1,1,1))
            #Lnew = torch.sum(mass * torch.cross(coordinates, velocities), dim=1)
            #_, Tnew = self.kinetic_energy(const, mass, species, velocities)
            #print("After remove angular momentum: ", Lnew, T0-Tnew)
            del r_com, v_com, L, I, omega, T0, T1, alpha
//  [python]
def zero_com(self, const, species, coordinates, velocities):
        """
        remove center of mass
        remove center of mass velocity
        remove angular momentum
        scale temperature back
        mass shape (nmol, molsize, 1)
        velocities shape (nmol, molsize, 3)
        """
        #mass for padding atom is 0.0
        mass = const.mass[species].unsqueeze(2)
        M = torch.sum(mass,dim=1,keepdim=True)
        _, T0 = self.kinetic_energy(const, mass, species, velocities)
        with torch.no_grad():
            r_com = torch.sum(mass*coordinates,dim=1,keepdim=True)/M
            coordinates.sub_(r_com)
            v_com = torch.sum(mass*velocities,dim=1,keepdim=True)/M
            velocities.sub_(v_com)
            L = torch.sum(mass * torch.cross(coordinates, velocities, dim=2), dim=1)
            I = torch.sum(mass * torch.norm(coordinates, dim=2, keepdim=True)**2, dim=1, keepdim=True) \
                * torch.eye(3, dtype=coordinates.dtype, device=coordinates.device).reshape(1,3,3) \
                - torch.sum( mass.unsqueeze(3) * coordinates.unsqueeze(3) * coordinates.unsqueeze(2), dim=1)
            # omega, _ = torch.solve(L.unsqueeze(2), I)
            omega = torch.linalg.solve(I, L.unsqueeze(2))
            velocities.add_( torch.cross(coordinates, omega.reshape(-1,1,3).repeat(1,coordinates.shape[1],1)) )
            _, T1 = self.kinetic_energy(const, mass, species, velocities)
            alpha = torch.sqrt(T0/T1)
            velocities.mul_(alpha.reshape(-1,1,1))
            #Lnew = torch.sum(mass * torch.cross(coordinates, velocities), dim=1)
            #_, Tnew = self.kinetic_energy(const, mass, species, velocities)
            #print("After remove angular momentum: ", Lnew, T0-Tnew)
            del r_com, v_com, L, I, omega, T0, T1, alpha
//  [python]
def raytrace_triangle(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangle: Float[Tensor, "trianglePoints=3 dims=3"]
) -> Bool[Tensor, "nrays"]:
    '''
    For each ray, return True if the triangle intersects that ray.
    '''
    NR = rays.size(0)

    # Triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    A, B, C = einops.repeat(triangle, "pts dims -> pts NR dims", NR=NR)
    assert A.shape == (NR, 3)

    # Each element of `rays` is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    O, D = rays.unbind(dim=1)
    assert O.shape == (NR, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
    
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    # Note - this works because mat[is_singular] has shape (NR_where_singular, 3, 3), so we
    # can broadcast the identity matrix to that shape.
    dets: Float[Tensor, "NR"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-8
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec = O - A

    # Solve eqns
    sol: Float[Tensor, "NR 3"] = t.solve(mat, vec)
    s, u, v = sol.unbind(dim=-1)

    # Return boolean of (matrix is nonsingular, and solution is in correct range implying intersection)
    return ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
//  [python]
def raytrace_triangle(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangle: Float[Tensor, "trianglePoints=3 dims=3"]
) -> Bool[Tensor, "nrays"]:
    '''
    For each ray, return True if the triangle intersects that ray.
    '''
    NR = rays.size(0)

    # Triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    A, B, C = einops.repeat(triangle, "pts dims -> pts NR dims", NR=NR)
    assert A.shape == (NR, 3)

    # Each element of `rays` is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    O, D = rays.unbind(dim=1)
    assert O.shape == (NR, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
    
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    # Note - this works because mat[is_singular] has shape (NR_where_singular, 3, 3), so we
    # can broadcast the identity matrix to that shape.
    dets: Float[Tensor, "NR"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-8
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec = O - A

    # Solve eqns
    sol: Float[Tensor, "NR 3"] = t.linalg.solve(mat, vec)
    s, u, v = sol.unbind(dim=-1)

    # Return boolean of (matrix is nonsingular, and solution is in correct range implying intersection)
    return ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
//  [python]
def mult_std_inv(std, x):
        return torch.solve(x, std)[1]
//  [python]
def mult_std_inv(std, x):
        return torch.linalg.solve(std, x)
//  [python]
def pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    return torch.solve(p_sqrt.float(),q_sqrt.float())[0].half()
//  [python]
def pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    return torch.linalg.solve(q_sqrt.float(),p_sqrt.float()).half()
//  [python]
def raytrace_mesh(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"]
) -> Float[Tensor, "nrays"]:
    '''
    For each ray, return the distance to the closest intersecting triangle, or infinity.
    '''
    NR = rays.size(0)
    NT = triangles.size(0)

    # Each triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    triangles = einops.repeat(triangles, "NT pts dims -> pts NR NT dims", NR=NR)
    A, B, C = triangles
    assert A.shape == (NR, NT, 3)

    # Each ray is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    rays = einops.repeat(rays, "NR pts dims -> pts NR NT dims", NT=NT)
    O, D = rays
    assert O.shape == (NR, NT, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR NT 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    dets: Float[Tensor, "NR NT"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-8
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec: Float[Tensor, "NR NT 3"] = O - A

    # Solve eqns (note, s is the distance along ray)
    sol: Float[Tensor, "NR NT 3"] = t.solve(vec, mat)
    s, u, v = sol.unbind(-1)

    # Get boolean of intersects, and use it to set distance to infinity wherever there is no intersection
    intersects = ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
    s[~intersects] = float("inf") # t.inf

    # Get the minimum distance (over all triangles) for each ray
    return s.min(dim=-1).values
//  [python]
def raytrace_mesh(
    rays: Float[Tensor, "nrays rayPoints=2 dims=3"],
    triangles: Float[Tensor, "ntriangles trianglePoints=3 dims=3"]
) -> Float[Tensor, "nrays"]:
    '''
    For each ray, return the distance to the closest intersecting triangle, or infinity.
    '''
    NR = rays.size(0)
    NT = triangles.size(0)

    # Each triangle is [[Ax, Ay, Az], [Bx, By, Bz], [Cx, Cy, Cz]]
    triangles = einops.repeat(triangles, "NT pts dims -> pts NR NT dims", NR=NR)
    A, B, C = triangles
    assert A.shape == (NR, NT, 3)

    # Each ray is [[Ox, Oy, Oz], [Dx, Dy, Dz]]
    rays = einops.repeat(rays, "NR pts dims -> pts NR NT dims", NT=NT)
    O, D = rays
    assert O.shape == (NR, NT, 3)

    # Define matrix on left hand side of equation
    mat: Float[Tensor, "NR NT 3 3"] = t.stack([- D, B - A, C - A], dim=-1)
    # Get boolean of where matrix is singular, and replace it with the identity in these positions
    dets: Float[Tensor, "NR NT"] = t.linalg.det(mat)
    is_singular = dets.abs() < 1e-8
    mat[is_singular] = t.eye(3)

    # Define vector on the right hand side of equation
    vec: Float[Tensor, "NR NT 3"] = O - A

    # Solve eqns (note, s is the distance along ray)
    sol: Float[Tensor, "NR NT 3"] = t.linalg.solve(mat, vec)
    s, u, v = sol.unbind(-1)

    # Get boolean of intersects, and use it to set distance to infinity wherever there is no intersection
    intersects = ((u >= 0) & (v >= 0) & (u + v <= 1) & ~is_singular)
    s[~intersects] = float("inf") # t.inf

    # Get the minimum distance (over all triangles) for each ray
    return s.min(dim=-1).values
//  [python]
def _kalman_gain(cov: Tensor, H: Tensor, R: Tensor) -> Tensor:
        Ht = H.permute(0, 2, 1)  # (1, 3, 3)
        covs_measured = cov @ Ht
        # Hcov = torch.matmul(H, cov)
        # print(Hcov.shape) # (N, 3, 3)
        system_covariance = torch.baddbmm(R, H@cov, Ht.repeat(cov.size(0),1,1))
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt = torch.solve(B, A)[1]
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def _kalman_gain(cov: Tensor, H: Tensor, R: Tensor) -> Tensor:
        Ht = H.permute(0, 2, 1)  # (1, 3, 3)
        covs_measured = cov @ Ht
        # Hcov = torch.matmul(H, cov)
        # print(Hcov.shape) # (N, 3, 3)
        system_covariance = torch.baddbmm(R, H@cov, Ht.repeat(cov.size(0),1,1))
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt = torch.linalg.solve(A, B)
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def mult_std_inv(std, x):
        ndim = x.shape[1]
        if ndim > 2:
            std = std.permute([0, 3, 1, 2])
            x = x.permute([0, 2, 1])
            sol = torch.solve(x, std)[0]
            sol = sol.permute([0, 2, 1]).contiguous()
        elif ndim == 2:
            # manual 2x2 system solution
            # | a b | | y1 |   | x_1 |
            # | c d | | y2 | = | x_2 |
            a = std[:, 0, 0, :]
            b = std[:, 0, 1, :]
            c = std[:, 1, 0, :]
            d = std[:, 1, 1, :]
            x1 = x[:, 0, :]
            x2 = x[:, 1, :]
            div = torch.reciprocal(a * d - c * b)
            y1 = div * (d * x1 - b * x2)
            y2 = div * (a * x2 - c * x1)
            sol = torch.stack((y1, y2), dim=1)
        else:
            raise ValueError(f"ndim={ndim}, it should be >= 2")
        return sol
//  [python]
def mult_std_inv(std, x):
        ndim = x.shape[1]
        if ndim > 2:
            std = std.permute([0, 3, 1, 2])
            x = x.permute([0, 2, 1])
            sol = torch.linalg.solve(std, x)
            sol = sol.permute([0, 2, 1]).contiguous()
        elif ndim == 2:
            # manual 2x2 system solution
            # | a b | | y1 |   | x_1 |
            # | c d | | y2 | = | x_2 |
            a = std[:, 0, 0, :]
            b = std[:, 0, 1, :]
            c = std[:, 1, 0, :]
            d = std[:, 1, 1, :]
            x1 = x[:, 0, :]
            x2 = x[:, 1, :]
            div = torch.reciprocal(a * d - c * b)
            y1 = div * (d * x1 - b * x2)
            y2 = div * (a * x2 - c * x1)
            sol = torch.stack((y1, y2), dim=1)
        else:
            raise ValueError(f"ndim={ndim}, it should be >= 2")
        return sol
//  [python]
def optimise_ridge_parameter(self,Features,Y):
        ridges=10.0**np.arange(-8,9)
        num_val_samples=int(Features.shape[0]*0.8)
        losses=[]
        Q_val=Features[0:num_val_samples,:].T @ Y[0:num_val_samples,:]
        G_val=Features[0:num_val_samples,:].T @ Features[0:num_val_samples,:]
        for ridge in ridges:
            Wo=torch.solve(Q_val,G_val+ridge*torch.eye(G_val.size(dim=0)))[0].T #better nmerical stability than .inv
            Y_train_pred=Features[num_val_samples::,:]@Wo.T
            losses.append(F.mse_loss(Y_train_pred,Y[num_val_samples::,:]))
        ridge=ridges[np.argmin(np.array(losses))]
        logging.info("Optimal lambda: "+str(ridge))
        return ridge
//  [python]
def optimise_ridge_parameter(self,Features,Y):
        ridges=10.0**np.arange(-8,9)
        num_val_samples=int(Features.shape[0]*0.8)
        losses=[]
        Q_val=Features[0:num_val_samples,:].T @ Y[0:num_val_samples,:]
        G_val=Features[0:num_val_samples,:].T @ Features[0:num_val_samples,:]
        for ridge in ridges:
            Wo=torch.linalg.solve(G_val+ridge*torch.eye(G_val.size(dim=0)),Q_val).T #better nmerical stability than .inv
            Y_train_pred=Features[num_val_samples::,:]@Wo.T
            losses.append(F.mse_loss(Y_train_pred,Y[num_val_samples::,:]))
        ridge=ridges[np.argmin(np.array(losses))]
        logging.info("Optimal lambda: "+str(ridge))
        return ridge
//  [python]
def get_result_of_mmse(h, p) -> TEN:  # MMSE beamformer
        h = h[0] + 1j * h[1]
        k, n = h.shape
        eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
        w = th.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))[0]
        w = w / (th.norm(w, dim=0, keepdim=True) * k ** 0.5)  # w.shape == [K, N]
        return th.stack((w.real, w.imag), dim=0)
//  [python]
def get_result_of_mmse(h, p) -> TEN:  # MMSE beamformer
        h = h[0] + 1j * h[1]
        k, n = h.shape
        eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
        w = th.linalg.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))
        w = w / (th.norm(w, dim=0, keepdim=True) * k ** 0.5)  # w.shape == [K, N]
        return th.stack((w.real, w.imag), dim=0)
//  [python]
def solve_grf(mass_mat,
              desired_acc,
              base_rot_mat_t,
              Wq,
              Wf: float,
              foot_friction_coef: float,
              clip_grf: bool,
              foot_contact_state,
              device: str = 'cuda'):
  num_envs = mass_mat.shape[0]
  g = torch.zeros((num_envs, 6), device=device)
  g[:, 2] = 9.8

  g[:, :3] = torch.matmul(base_rot_mat_t, g[:, :3, None])[:, :, 0]
  Q = torch.zeros((num_envs, 6, 6), device=device) + Wq[None, :]
  Wf_mat = torch.eye(12, device=device) * Wf
  R = torch.zeros((num_envs, 12, 12), device=device) + Wf_mat[None, :]

  quad_term = torch.bmm(torch.bmm(torch.transpose(mass_mat, 1, 2), Q),
                        mass_mat) + R
  linear_term = torch.bmm(torch.bmm(torch.transpose(mass_mat, 1, 2), Q),
                          (g + desired_acc)[:, :, None])[:, :, 0]
  grf = torch.solve(linear_term, quad_term)[0]

  base_rot_mat = torch.transpose(base_rot_mat_t, 1, 2)
  grf = grf.reshape((-1, 4, 3))
  grf_world = torch.transpose(
      torch.bmm(base_rot_mat, torch.transpose(grf, 1, 2)), 1, 2)
  if clip_grf:
    grf_world[:, :, 2] = grf_world[:, :, 2].clip(min=10, max=130)
    grf_world[:, :, 2] *= foot_contact_state
  friction_force = torch.norm(grf_world[:, :, :2], dim=2) + 0.001
  max_friction_force = foot_friction_coef * grf_world[:, :, 2].clip(min=0)
  multiplier = torch.where(friction_force < max_friction_force, 1,
                           max_friction_force / friction_force)
  if clip_grf:
    grf_world[:, :, :2] *= multiplier[:, :, None]
  grf = torch.transpose(
      torch.bmm(base_rot_mat_t, torch.transpose(grf_world, 1, 2)), 1, 2)
  grf = grf.reshape((-1, 12))

  # Convert to motor torques
  solved_acc = torch.bmm(mass_mat, grf[:, :, None])[:, :, 0] - g
  qp_cost = torch.bmm(
      torch.bmm((solved_acc - desired_acc)[:, :, None].transpose(1, 2), Q),
      (solved_acc - desired_acc)[:, :, None])[:, 0, 0]
  return grf, solved_acc, qp_cost, torch.sum(
      friction_force > max_friction_force + 1, dim=1)
//  [python]
def solve_grf(mass_mat,
              desired_acc,
              base_rot_mat_t,
              Wq,
              Wf: float,
              foot_friction_coef: float,
              clip_grf: bool,
              foot_contact_state,
              device: str = 'cuda'):
  num_envs = mass_mat.shape[0]
  g = torch.zeros((num_envs, 6), device=device)
  g[:, 2] = 9.8

  g[:, :3] = torch.matmul(base_rot_mat_t, g[:, :3, None])[:, :, 0]
  Q = torch.zeros((num_envs, 6, 6), device=device) + Wq[None, :]
  Wf_mat = torch.eye(12, device=device) * Wf
  R = torch.zeros((num_envs, 12, 12), device=device) + Wf_mat[None, :]

  quad_term = torch.bmm(torch.bmm(torch.transpose(mass_mat, 1, 2), Q),
                        mass_mat) + R
  linear_term = torch.bmm(torch.bmm(torch.transpose(mass_mat, 1, 2), Q),
                          (g + desired_acc)[:, :, None])[:, :, 0]
  grf = torch.linalg.solve(quad_term, linear_term)

  base_rot_mat = torch.transpose(base_rot_mat_t, 1, 2)
  grf = grf.reshape((-1, 4, 3))
  grf_world = torch.transpose(
      torch.bmm(base_rot_mat, torch.transpose(grf, 1, 2)), 1, 2)
  if clip_grf:
    grf_world[:, :, 2] = grf_world[:, :, 2].clip(min=10, max=130)
    grf_world[:, :, 2] *= foot_contact_state
  friction_force = torch.norm(grf_world[:, :, :2], dim=2) + 0.001
  max_friction_force = foot_friction_coef * grf_world[:, :, 2].clip(min=0)
  multiplier = torch.where(friction_force < max_friction_force, 1,
                           max_friction_force / friction_force)
  if clip_grf:
    grf_world[:, :, :2] *= multiplier[:, :, None]
  grf = torch.transpose(
      torch.bmm(base_rot_mat_t, torch.transpose(grf_world, 1, 2)), 1, 2)
  grf = grf.reshape((-1, 12))

  # Convert to motor torques
  solved_acc = torch.bmm(mass_mat, grf[:, :, None])[:, :, 0] - g
  qp_cost = torch.bmm(
      torch.bmm((solved_acc - desired_acc)[:, :, None].transpose(1, 2), Q),
      (solved_acc - desired_acc)[:, :, None])[:, 0, 0]
  return grf, solved_acc, qp_cost, torch.sum(
      friction_force > max_friction_force + 1, dim=1)
//  [python]
def _analytic_scale_optim_unnormalized(
    K: torch.ndarray,
    nn_targets: torch.ndarray,
) -> torch.ndarray:
    return torch.sum(
        torch.einsum(
            "ijk,ijk->ik", nn_targets, torch.solve(nn_targets, K)[0]
        ),
        axis=0,
    )
//  [python]
def _analytic_scale_optim_unnormalized(
    K: torch.ndarray,
    nn_targets: torch.ndarray,
) -> torch.ndarray:
    return torch.sum(
        torch.einsum(
            "ijk,ijk->ik", nn_targets, torch.linalg.solve(K, nn_targets)
        ),
        axis=0,
    )
//  [python]
def get_Phi_batched(X, X_is, radius):
    dist_vec = get_weights_batched(X, X_is, radius=radius)  # (200,)
    M = get_moment_mat_batched(X_is, dist_vec)  # (4,4)
    p_x = torch.cat((torch.ones((X.shape[0], X.shape[1], 1)).to(X.device).to(X.dtype), X), dim=-1)
    P_Xis = torch.cat((torch.ones((X_is.shape[0], X_is.shape[1], 1)).to(X_is.device).to(X_is.dtype),
                       X_is, ),dim=-1,)

    M_inv = torch.solve(eye_like(M), M)[0]

    first_term = torch.einsum('bij,bijk->bik', p_x, M_inv)
    second_term = torch.einsum('bpn,bmn->bmpn', P_Xis.transpose(2, 1), dist_vec)
    Phi = torch.einsum('bij,bijk->bik', first_term, second_term)
    return Phi
//  [python]
def get_Phi_batched(X, X_is, radius):
    dist_vec = get_weights_batched(X, X_is, radius=radius)  # (200,)
    M = get_moment_mat_batched(X_is, dist_vec)  # (4,4)
    p_x = torch.cat((torch.ones((X.shape[0], X.shape[1], 1)).to(X.device).to(X.dtype), X), dim=-1)
    P_Xis = torch.cat((torch.ones((X_is.shape[0], X_is.shape[1], 1)).to(X_is.device).to(X_is.dtype),
                       X_is, ),dim=-1,)

    M_inv = torch.linalg.solve(M, eye_like(M))

    first_term = torch.einsum('bij,bijk->bik', p_x, M_inv)
    second_term = torch.einsum('bpn,bmn->bmpn', P_Xis.transpose(2, 1), dist_vec)
    Phi = torch.einsum('bij,bijk->bik', first_term, second_term)
    return Phi
//  [python]
def exactsolve(A: LinearOperator, B: torch.Tensor, E: Union[torch.Tensor, None],
               M: Union[LinearOperator, None]):
    """
    Solve the linear equation by contructing the full matrix of LinearOperators.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.differentiation_utils import LinearOperator
    >>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
    >>> B = torch.tensor([[5., 6], [7, 8]])
    >>> exactsolve(A, B, None, None)
    tensor([[-3., -4.],
            [ 4.,  5.]])

    Parameters
    ----------
    A: LinearOperator
        The linear operator A to be solved. Shape: (`*BA`, na, na)
    B: torch.Tensor
        Batched matrix B. Shape: (`*BB`, na, ncols)
    E: torch.Tensor or None
        Batched vector E. Shape: (`*BE`, ncols)
    M: LinearOperator or None
        The linear operator M. Shape: (`*BM`, na, na)

    Returns
    -------
    torch.Tensor
        The Solution matrix X. Shape: (`*BBE`, na, ncols)

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.

    """
    if E is None:
        Amatrix = A.fullmatrix()
        x = torch.solve(B, Amatrix)[0]
    elif M is None:
        Amatrix = A.fullmatrix()
        x = solve_ABE(Amatrix, B, E)
    else:
        Mmatrix = M.fullmatrix()
        L = torch.linalg.cholesky(Mmatrix)
        Linv = torch.inverse(L)
        LinvT = Linv.transpose(-2, -1).conj()
        A2 = torch.matmul(Linv, A.mm(LinvT))
        B2 = torch.matmul(Linv, B)

        X2 = solve_ABE(A2, B2, E)
        x = torch.matmul(LinvT, X2)
    return x
//  [python]
def exactsolve(A: LinearOperator, B: torch.Tensor, E: Union[torch.Tensor, None],
               M: Union[LinearOperator, None]):
    """
    Solve the linear equation by contructing the full matrix of LinearOperators.

    Examples
    --------
    >>> import torch
    >>> from deepchem.utils.differentiation_utils import LinearOperator
    >>> A = LinearOperator.m(torch.tensor([[1., 2], [3, 4]]))
    >>> B = torch.tensor([[5., 6], [7, 8]])
    >>> exactsolve(A, B, None, None)
    tensor([[-3., -4.],
            [ 4.,  5.]])

    Parameters
    ----------
    A: LinearOperator
        The linear operator A to be solved. Shape: (`*BA`, na, na)
    B: torch.Tensor
        Batched matrix B. Shape: (`*BB`, na, ncols)
    E: torch.Tensor or None
        Batched vector E. Shape: (`*BE`, ncols)
    M: LinearOperator or None
        The linear operator M. Shape: (`*BM`, na, na)

    Returns
    -------
    torch.Tensor
        The Solution matrix X. Shape: (`*BBE`, na, ncols)

    Warnings
    --------
    * As this method construct the linear operators explicitly, it might requires
      a large memory.

    """
    if E is None:
        Amatrix = A.fullmatrix()
        x = torch.linalg.solve(Amatrix, B)
    elif M is None:
        Amatrix = A.fullmatrix()
        x = solve_ABE(Amatrix, B, E)
    else:
        Mmatrix = M.fullmatrix()
        L = torch.linalg.cholesky(Mmatrix)
        Linv = torch.inverse(L)
        LinvT = Linv.transpose(-2, -1).conj()
        A2 = torch.matmul(Linv, A.mm(LinvT))
        B2 = torch.matmul(Linv, B)

        X2 = solve_ABE(A2, B2, E)
        x = torch.matmul(LinvT, X2)
    return x
//  [python]
def replace_fc(self,trainloader):
        self._network = self._network.eval()

        if self.args['use_RP']:
            #these lines are needed because the CosineLinear head gets deleted between streams and replaced by one with more classes (for CIL)
            self._network.fc.use_RP=True
            if self.args['M']>0:
                self._network.fc.W_rand=self.W_rand
            else:
                self._network.fc.W_rand=None

        Features_f = []
        label_list = []
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (_,data,label)=batch
                data=data.cuda()
                label=label.cuda()
                embedding = self._network.convnet(data)
                Features_f.append(embedding.cpu())
                label_list.append(label.cpu())
        Features_f = torch.cat(Features_f, dim=0)
        label_list = torch.cat(label_list, dim=0)
        
        Y=target2onehot(label_list,self.total_classnum)
        if self.args['use_RP']:
            #print('Number of pre-trained feature dimensions = ',Features_f.shape[-1])
            if self.args['M']>0:
                Features_h=torch.nn.functional.relu(Features_f@ self._network.fc.W_rand.cpu())
            else:
                Features_h=Features_f
            self.Q=self.Q+Features_h.T @ Y 
            self.G=self.G+Features_h.T @ Features_h
            ridge=self.optimise_ridge_parameter(Features_h,Y)
            Wo=torch.solve(self.Q,self.G+ridge*torch.eye(self.G.size(dim=0)))[0].T #better nmerical stability than .inv
            self._network.fc.weight.data=Wo[0:self._network.fc.weight.shape[0],:].to(device='cuda')
        else:
            for class_index in np.unique(self.train_dataset.labels):
                data_index=(label_list==class_index).nonzero().squeeze(-1)
                if self.is_dil:
                    class_prototype=Features_f[data_index].sum(0)
                    self._network.fc.weight.data[class_index]+=class_prototype.to(device='cuda') #for dil, we update all classes in all tasks
                else:
                    #original cosine similarity approach of Zhou et al (2023)
                    class_prototype=Features_f[data_index].mean(0)
                    self._network.fc.weight.data[class_index]=class_prototype
//  [python]
def replace_fc(self,trainloader):
        self._network = self._network.eval()

        if self.args['use_RP']:
            #these lines are needed because the CosineLinear head gets deleted between streams and replaced by one with more classes (for CIL)
            self._network.fc.use_RP=True
            if self.args['M']>0:
                self._network.fc.W_rand=self.W_rand
            else:
                self._network.fc.W_rand=None

        Features_f = []
        label_list = []
        with torch.no_grad():
            for i, batch in enumerate(trainloader):
                (_,data,label)=batch
                data=data.cuda()
                label=label.cuda()
                embedding = self._network.convnet(data)
                Features_f.append(embedding.cpu())
                label_list.append(label.cpu())
        Features_f = torch.cat(Features_f, dim=0)
        label_list = torch.cat(label_list, dim=0)
        
        Y=target2onehot(label_list,self.total_classnum)
        if self.args['use_RP']:
            #print('Number of pre-trained feature dimensions = ',Features_f.shape[-1])
            if self.args['M']>0:
                Features_h=torch.nn.functional.relu(Features_f@ self._network.fc.W_rand.cpu())
            else:
                Features_h=Features_f
            self.Q=self.Q+Features_h.T @ Y 
            self.G=self.G+Features_h.T @ Features_h
            ridge=self.optimise_ridge_parameter(Features_h,Y)
            Wo=torch.linalg.solve(self.G+ridge*torch.eye(self.G.size(dim=0)),self.Q).T #better nmerical stability than .inv
            self._network.fc.weight.data=Wo[0:self._network.fc.weight.shape[0],:].to(device='cuda')
        else:
            for class_index in np.unique(self.train_dataset.labels):
                data_index=(label_list==class_index).nonzero().squeeze(-1)
                if self.is_dil:
                    class_prototype=Features_f[data_index].sum(0)
                    self._network.fc.weight.data[class_index]+=class_prototype.to(device='cuda') #for dil, we update all classes in all tasks
                else:
                    #original cosine similarity approach of Zhou et al (2023)
                    class_prototype=Features_f[data_index].mean(0)
                    self._network.fc.weight.data[class_index]=class_prototype
//  [python]
def batch_inv(x: torch.Tensor):
    """batch inverse
    using torch.linalg.solve to accelerate (computation on GPU)
    Args:
        x (torch.Tensor): (B,M,M)
    Returns:
        inv_x (torch.Tensor): (B,M,M)
    """
    assert x.size(1) == x.size(2), "batch inv can only compute inverse of square matrix, get size:{}".format(x.size())
    M = x.size(1)  
    valid = torch.det(x)!=0
    x_ = x[valid]
    inv_x = torch.zeros_like(x)
    # if torch.abs(x).min().item() == 0:
    #     x += 1e-6  # jitter for inv operation
    identity = torch.eye(M)[None,...].to(x.device)  # (1,M,N)
    inv_x[valid] = torch.inverse(x_)  # need torch>=1.8
    return inv_x
//  [python]
def batch_inv(x: torch.Tensor):
    """batch inverse
    using torch.linalg.solve to accelerate (computation on GPU)
    Args:
        x (torch.Tensor): (B,M,M)
    Returns:
        inv_x (torch.Tensor): (B,M,M)
    """
    assert x.size(1) == x.size(2), "batch inv can only compute inverse of square matrix, get size:{}".format(x.size())
    M = x.size(1)  
    valid = torch.det(x)!=0
    x_ = x[valid]
    inv_x = torch.zeros_like(x)
    # if torch.abs(x).min().item() == 0:
    #     x += 1e-6  # jitter for inv operation
    identity = torch.eye(M)[None,...].to(x.device)  # (1,M,N)
    inv_x[valid] = torch.linalg.solve(x_,identity)  # need torch>=1.8
    return inv_x
//  [python]
def fit(self, features: torch.Tensor, targets: torch.Tensor, prior_precision: float = 1):
        """
        Learn weights from data using MAP inference
        Args:
            features: (Npoints x Nfeatures) tensor with training features
            targets: (Npoints,) tensor
            prior_precision: Precision of an isotropic Gaussian prior (also known as ridge regulariser)
        Returns:
            None
        """
        assert len(targets) == features.shape[0]
        assert len(features.shape) == 2

        self.posterior_prec = (
            features.T @ features
            + torch.eye(features.shape[1], dtype=features.dtype, device=features.device) * prior_precision
        )

        self.w = torch.solve(features.T @ targets, self.posterior_prec)[0]
//  [python]
def fit(self):
        if self.readout_training in {'gd', 'svd'}:
            return

        if self.readout_training == 'cholesky':
            W = torch.solve(self.XTy,
                                   self.XTX + self.lambda_reg * torch.eye(
                                       self.XTX.size(0), device=self.XTX.device))[0].t()
            self.XTX = None
            self.XTy = None

            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
        elif self.readout_training == 'inv':
            I = (self.lambda_reg * torch.eye(self.XTX.size(0))).to(
                self.XTX.device)
            A = self.XTX + I
            X_rank = torch.linalg.matrix_rank(A).item()

            if X_rank == self.X.size(0):
                W = torch.mm(torch.inverse(A), self.XTy).t()
            else:
                W = torch.mm(torch.pinverse(A), self.XTy).t()

            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])

            self.XTX = None
            self.XTy = None
//  [python]
def fit(self):
        if self.readout_training in {'gd', 'svd'}:
            return

        if self.readout_training == 'cholesky':
            W = torch.linalg.solve(self.XTy,
                                   self.XTX + self.lambda_reg * torch.eye(
                                       self.XTX.size(0), device=self.XTX.device))[0].t()
            self.XTX = None
            self.XTy = None

            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
        elif self.readout_training == 'inv':
            I = (self.lambda_reg * torch.eye(self.XTX.size(0))).to(
                self.XTX.device)
            A = self.XTX + I
            X_rank = torch.linalg.matrix_rank(A).item()

            if X_rank == self.X.size(0):
                W = torch.mm(torch.inverse(A), self.XTy).t()
            else:
                W = torch.mm(torch.pinverse(A), self.XTy).t()

            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])

            self.XTX = None
            self.XTy = None
//  [python]
def solve_reg(x_tr, y_tr, x_te, y_te, lam, y_std=1.0):
    n, d = x_tr.shape
    xtype = x_tr.type()

    if x_tr.shape[0] > x_tr.shape[1]:
        b = x_tr.T @ (y_tr[:, None] if y_tr.dim() == 1 else y_tr)
        y_pred = x_te @ torch.solve(b, x_tr.T @ x_tr + lam * torch.eye(d).type(xtype))[0]
    else:
        b = y_tr[:, None] if y_tr.dim() == 1 else y_tr
        y_pred = x_te @ (x_tr.T @ torch.solve(b, x_tr @ x_tr.T + lam * torch.eye(n).type(xtype))[0])

    mse = mean_squared_error((y_pred * y_std).cpu(), (y_te * y_std).cpu())
    acc = ((y_pred.argmax(axis=1) == y_te.argmax(axis=1)) * 1.0).mean().item()
    return mse, acc