jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        x_plus = torch.ones(len(x), 1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        #print('weighted_x_2d',weighted_x_2d) #PESI
        #print('y_actual_2d',y_actual_2d) #RAPPRESENTAZIONE ONE OUT DEL TARGET

        # Use gels to do LSE, then pick out the solution rows:
        try:
            #coeff_2d, _ = torch.gels(y_actual_2d, weighted_x_2d)
            coeff_2d, _ = torch.linalg.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)
//  [python]
def _test_acos_outside_bounds(self, x, y, dydx, bound):
        """
        Check that `acos_linear_extrapolation` yields points on a line with correct
        slope, and that the function is continuous around `bound`.
        """
        bound_t = torch.tensor(bound, device=x.device, dtype=x.dtype)
        # fit a line: slope * x + bias = y
        x_1 = torch.stack([x, torch.ones_like(x)], dim=-1)
        slope, bias = torch.lstsq(y[:, None], x_1).solution.view(-1)[:2]
        desired_slope = (-1.0) / torch.sqrt(1.0 - bound_t**2)
        # test that the desired slope is the same as the fitted one
        self.assertClose(desired_slope.view(1), slope.view(1), atol=1e-2)
        # test that the autograd's slope is the same as the desired one
        self.assertClose(desired_slope.expand_as(dydx), dydx, atol=1e-2)
        # test that the value of the fitted line at x=bound equals
        # arccos(x), i.e. the function is continuous around the bound
        y_bound_lin = (slope * bound_t + bias).view(1)
        y_bound_acos = bound_t.acos().view(1)
        self.assertClose(y_bound_lin, y_bound_acos, atol=1e-2)
//  [python]
def _test_acos_outside_bounds(self, x, y, dydx, bound):
        """
        Check that `acos_linear_extrapolation` yields points on a line with correct
        slope, and that the function is continuous around `bound`.
        """
        bound_t = torch.tensor(bound, device=x.device, dtype=x.dtype)
        # fit a line: slope * x + bias = y
        x_1 = torch.stack([x, torch.ones_like(x)], dim=-1)
        slope, bias = torch.linalg.lstsq(x_1, y[:, None]).solution.view(-1)[:2]
        desired_slope = (-1.0) / torch.sqrt(1.0 - bound_t**2)
        # test that the desired slope is the same as the fitted one
        self.assertClose(desired_slope.view(1), slope.view(1), atol=1e-2)
        # test that the autograd's slope is the same as the desired one
        self.assertClose(desired_slope.expand_as(dydx), dydx, atol=1e-2)
        # test that the value of the fitted line at x=bound equals
        # arccos(x), i.e. the function is continuous around the bound
        y_bound_lin = (slope * bound_t + bias).view(1)
        y_bound_acos = bound_t.acos().view(1)
        self.assertClose(y_bound_lin, y_bound_acos, atol=1e-2)
//  [python]
def _get_perspective_coeffs(startpoints: List[List[int]],
                            endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor(
            [p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor(
            [0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.solve(b_matrix.unsqueeze(1), a_matrix).solution

    output: List[float] = res.tolist()
    return output
//  [python]
def _get_perspective_coeffs(startpoints: List[List[int]],
                            endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor(
            [p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor(
            [0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.linalg.lstsq(a_matrix, b_matrix, driver='gels').solution

    output: List[float] = res.tolist()
    return output
//  [python]
def lstsq_with_handles(A, b, handle_idx, handle_pos):
    b = b - A[:, handle_idx] @ handle_pos
    handle_mask = torch.zeros_like(A[:, 0], dtype=bool)
    handle_mask[handle_idx] = 1
    L = A[:, handle_mask.logical_not()]
    x = torch.lstsq(L, b)[0]
    x_out = torch.zeros_like(b)
    x_out[handle_idx] = handle_pos
    x_out[handle_mask.logical_not()] = x
    return x_out
//  [python]
def lstsq_with_handles(A, b, handle_idx, handle_pos):
    b = b - A[:, handle_idx] @ handle_pos
    handle_mask = torch.zeros_like(A[:, 0], dtype=bool)
    handle_mask[handle_idx] = 1
    L = A[:, handle_mask.logical_not()]
    x = torch.linalg.lstsq(L, b)[0]
    x_out = torch.zeros_like(b)
    x_out[handle_idx] = handle_pos
    x_out[handle_mask.logical_not()] = x
    return x_out
//  [python]
def intersect_rays(
    a_origins: Float[Tensor, "*#batch dim"],
    a_directions: Float[Tensor, "*#batch dim"],
    b_origins: Float[Tensor, "*#batch dim"],
    b_directions: Float[Tensor, "*#batch dim"],
) -> Float[Tensor, "*batch dim"]:
    """Compute the least-squares intersection of rays. Uses the math from here:
    https://math.stackexchange.com/a/1762491/286022
    """

    # Broadcast and stack the tensors.
    a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors(
        a_origins, a_directions, b_origins, b_directions
    )
    origins = torch.stack((a_origins, b_origins), dim=-2)
    directions = torch.stack((a_directions, b_directions), dim=-2)

    # Compute n_i * n_i^T - eye(3) from the equation.
    n = einsum(directions, directions, "... n i, ... n j -> ... n i j")
    n = n - torch.eye(3, dtype=origins.dtype, device=origins.device)

    # Compute the left-hand side of the equation.
    lhs = reduce(n, "... n i j -> ... i j", "sum")

    # Compute the right-hand side of the equation.
    rhs = einsum(n, origins, "... n i j, ... n j -> ... n i")
    rhs = reduce(rhs, "... n i -> ... i", "sum")

    # Left-matrix-multiply both sides by the inverse of lhs to find p.
    return torch.solve(rhs.unsqueeze(-1), lhs).solution.squeeze(-1)
//  [python]
def intersect_rays(
    a_origins: Float[Tensor, "*#batch dim"],
    a_directions: Float[Tensor, "*#batch dim"],
    b_origins: Float[Tensor, "*#batch dim"],
    b_directions: Float[Tensor, "*#batch dim"],
) -> Float[Tensor, "*batch dim"]:
    """Compute the least-squares intersection of rays. Uses the math from here:
    https://math.stackexchange.com/a/1762491/286022
    """

    # Broadcast and stack the tensors.
    a_origins, a_directions, b_origins, b_directions = torch.broadcast_tensors(
        a_origins, a_directions, b_origins, b_directions
    )
    origins = torch.stack((a_origins, b_origins), dim=-2)
    directions = torch.stack((a_directions, b_directions), dim=-2)

    # Compute n_i * n_i^T - eye(3) from the equation.
    n = einsum(directions, directions, "... n i, ... n j -> ... n i j")
    n = n - torch.eye(3, dtype=origins.dtype, device=origins.device)

    # Compute the left-hand side of the equation.
    lhs = reduce(n, "... n i j -> ... i j", "sum")

    # Compute the right-hand side of the equation.
    rhs = einsum(n, origins, "... n i j, ... n j -> ... n i")
    rhs = reduce(rhs, "... n i -> ... i", "sum")

    # Left-matrix-multiply both sides by the inverse of lhs to find p.
    return torch.linalg.lstsq(lhs, rhs).solution
//  [python]
def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3=None):
        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) * self.qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
        B_reg = (self.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=None):
        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) * self.qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
        B_reg = (self.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 _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.
    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )
    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.
    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.solve(b_matrix.unsqueeze(1), a_matrix).solution

    output: List[float] = res.tolist()
    return output
//  [python]
def _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.
    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )
    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.
    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.linalg.lstsq(a_matrix, b_matrix, driver="gels").solution

    output: List[float] = res.tolist()
    return output
//  [python]
def _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.solve(b_matrix.unsqueeze(1), a_matrix).solution

    output: List[float] = res.tolist()
    return output
//  [python]
def _get_perspective_coeffs(startpoints: List[List[int]], endpoints: List[List[int]]) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.linalg.lstsq(a_matrix, b_matrix, driver="gels").solution

    output: List[float] = res.tolist()
    return output
//  [python]
def vamp_score(X, Y, schatten_norm: int = 2, center_covariances: bool = True):
    """Variational Approach for learning Markov Processes (VAMP) score by :footcite:t:`Wu2019`.

    Args:
        X (torch.Tensor): Covariates for the initial time steps.
        Y (torch.Tensor): Covariates for the evolved time steps.
        schatten_norm (int, optional): Computes the VAMP-p score with ``p = schatten_norm``. Defaults to 2.
        center_covariances (bool, optional): Use centered covariances to compute the VAMP score. Defaults to True.

    Raises:
        NotImplementedError: If ``schatten_norm`` is not 1 or 2.

    """
    cov_X, cov_Y, cov_XY = (
        covariance(X, center=center_covariances),
        covariance(Y, center=center_covariances),
        covariance(X, Y, center=center_covariances),
    )
    if schatten_norm == 2:
        # Using least squares in place of pinv for numerical stability
        M_X = torch.solve(cov_XY, cov_X)[0]
        M_Y = torch.solve(cov_XY.T, cov_Y)[0]
        return torch.trace(M_X @ M_Y)
    elif schatten_norm == 1:
        sqrt_cov_X = sqrtmh(cov_X)
        sqrt_cov_Y = sqrtmh(cov_Y)
        M = torch.linalg.multi_dot(
            [
                torch.linalg.pinv(sqrt_cov_X, hermitian=True),
                cov_XY,
                torch.linalg.pinv(sqrt_cov_Y, hermitian=True),
            ]
        )
        return torch.linalg.matrix_norm(M, "nuc")
    else:
        raise NotImplementedError(f"Schatten norm {schatten_norm} not implemented")
//  [python]
def vamp_score(X, Y, schatten_norm: int = 2, center_covariances: bool = True):
    """Variational Approach for learning Markov Processes (VAMP) score by :footcite:t:`Wu2019`.

    Args:
        X (torch.Tensor): Covariates for the initial time steps.
        Y (torch.Tensor): Covariates for the evolved time steps.
        schatten_norm (int, optional): Computes the VAMP-p score with ``p = schatten_norm``. Defaults to 2.
        center_covariances (bool, optional): Use centered covariances to compute the VAMP score. Defaults to True.

    Raises:
        NotImplementedError: If ``schatten_norm`` is not 1 or 2.

    """
    cov_X, cov_Y, cov_XY = (
        covariance(X, center=center_covariances),
        covariance(Y, center=center_covariances),
        covariance(X, Y, center=center_covariances),
    )
    if schatten_norm == 2:
        # Using least squares in place of pinv for numerical stability
        M_X = torch.linalg.lstsq(cov_X, cov_XY).solution
        M_Y = torch.linalg.lstsq(cov_Y, cov_XY.T).solution
        return torch.trace(M_X @ M_Y)
    elif schatten_norm == 1:
        sqrt_cov_X = sqrtmh(cov_X)
        sqrt_cov_Y = sqrtmh(cov_Y)
        M = torch.linalg.multi_dot(
            [
                torch.linalg.pinv(sqrt_cov_X, hermitian=True),
                cov_XY,
                torch.linalg.pinv(sqrt_cov_Y, hermitian=True),
            ]
        )
        return torch.linalg.matrix_norm(M, "nuc")
    else:
        raise NotImplementedError(f"Schatten norm {schatten_norm} not implemented")
//  [python]
def estimate_amplitude(z, N, constrain=False, representation="fft"):
    V = complex_oscillator(z, N=N, constrain=constrain).sum(dim=-2)

    n = torch.arange(N, device=z.device)
    omega = torch.angle(z)

    U = torch.cos(omega[..., None, :] * n[None, :, None])

    if representation == "fft":
        U = torch.fft.rfft(U, dim=-2).abs()
        V = torch.fft.rfft(V).abs()

    least_squares_soln = torch.solve(U.t() @ U, U.t() @ V).solution
    return least_squares_soln
//  [python]
def estimate_amplitude(z, N, constrain=False, representation="fft"):
    V = complex_oscillator(z, N=N, constrain=constrain).sum(dim=-2)

    n = torch.arange(N, device=z.device)
    omega = torch.angle(z)

    U = torch.cos(omega[..., None, :] * n[None, :, None])

    if representation == "fft":
        U = torch.fft.rfft(U, dim=-2).abs()
        V = torch.fft.rfft(V).abs()

    least_squares_soln = torch.linalg.lstsq(U, V).solution
    return least_squares_soln
//  [python]
def least_squares(design, targets, trials=10, solver="gelsd"):
    """
    Minimizes
        || design @ solution - targets ||
    using torch.linalg.lstsq.

    Returns
        solution, predictions (= design @ solution)

    Since torch.linalg.lstsq is non-deteministic,
    the best solution of "trials" is return.

    As the design (kern \\otimes tasks_kern) matrix could fall into
    ill-conditioned form, xGELSY and xGELS fail often (for some reason - update the origin).

    - https://pytorch.org/docs/stable/generated/torch.linalg.lstsq.html#torch.linalg.lstsq
    - https://www.smcm.iqfr.csic.es/docs/intel/mkl/mkl_manual/lse/lse_drllsp.htm
    """
    best = None
    predictions = None
    for _ in range(trials):
        sol = torch.linalg.lstsq(design, targets, driver=solver)
        pred = design @ sol.solution
        err = (pred - targets).abs().mean()
        if not best or err < best:
            solution = sol.solution
            predictions = pred
            best = err
    return solution, predictions
//  [python]
def get_scale_aligned(
    predicted_quantity: Float[Tensor, "B N"],
    fixed_quantity: Float[Tensor, "B N"],
    mask: Float[Tensor, "B N"],
    fuse=True,
):
    """Solves for scale alignment in the overlapping region specified by mask.
    Correspondences are valid where mask==1.
    """
    device = predicted_quantity.device

    b, n = predicted_quantity.shape
    p = predicted_quantity.view(b, -1, 1)
    f = fixed_quantity.view(b, -1, 1)
    m = mask.view(b, -1)

    scale = []
    for i in range(b):
        # solve for A, where ||SX-B|. A is the scale
        X = p[i, m[i] == 1]
        B = f[i, m[i] == 1]
        S = torch.lstsq(B,X)[0]
        scale.append(S[0, 0])
    scale = torch.stack(scale).to(device)

    return S.view(b, 1) * predicted_quantity
//  [python]
def get_scale_aligned(
    predicted_quantity: Float[Tensor, "B N"],
    fixed_quantity: Float[Tensor, "B N"],
    mask: Float[Tensor, "B N"],
    fuse=True,
):
    """Solves for scale alignment in the overlapping region specified by mask.
    Correspondences are valid where mask==1.
    """
    device = predicted_quantity.device

    b, n = predicted_quantity.shape
    p = predicted_quantity.view(b, -1, 1)
    f = fixed_quantity.view(b, -1, 1)
    m = mask.view(b, -1)

    scale = []
    for i in range(b):
        # solve for A, where ||SX-B|. A is the scale
        X = p[i, m[i] == 1]
        B = f[i, m[i] == 1]
        S = torch.linalg.lstsq(X, B)[0]
        scale.append(S[0, 0])
    scale = torch.stack(scale).to(device)

    return S.view(b, 1) * predicted_quantity
//  [python]
def deepprojection_score(
    X,
    Y,
    relaxed: bool = True,
    metric_deformation: float = 1.0,
    center_covariances: bool = True,
):
    """Deep Projection score by :footcite:t:`Kostic2023DPNets`.

    Args:
        X (torch.Tensor): Covariates for the initial time steps.
        Y (torch.Tensor): Covariates for the evolved time steps.
        relaxed (bool, optional): Whether to use the relaxed (more numerically stable) or the full deep-projection loss. Defaults to True.
        metric_deformation (float, optional): Strength of the metric metric deformation loss: Defaults to 1.0.
        center_covariances (bool, optional): Use centered covariances to compute the VAMP score. Defaults to True.

    """
    cov_X, cov_Y, cov_XY = (
        covariance(X, center=center_covariances),
        covariance(Y, center=center_covariances),
        covariance(X, Y, center=center_covariances),
    )
    R_X = log_fro_metric_deformation_loss(cov_X)
    R_Y = log_fro_metric_deformation_loss(cov_Y)
    if relaxed:
        S = (torch.linalg.matrix_norm(cov_XY, ord="fro") ** 2) / (
            (
                torch.linalg.matrix_norm(cov_X, ord=2)
                * torch.linalg.matrix_norm(cov_Y, ord=2)
            )
        )
    else:
        M_X = torch.solve(cov_XY, cov_X).solution
        M_Y = torch.solve(cov_XY.T, cov_Y).solution
        S = torch.trace(M_X @ M_Y)
    return S - 0.5 * metric_deformation * (R_X + R_Y)
//  [python]
def deepprojection_score(
    X,
    Y,
    relaxed: bool = True,
    metric_deformation: float = 1.0,
    center_covariances: bool = True,
):
    """Deep Projection score by :footcite:t:`Kostic2023DPNets`.

    Args:
        X (torch.Tensor): Covariates for the initial time steps.
        Y (torch.Tensor): Covariates for the evolved time steps.
        relaxed (bool, optional): Whether to use the relaxed (more numerically stable) or the full deep-projection loss. Defaults to True.
        metric_deformation (float, optional): Strength of the metric metric deformation loss: Defaults to 1.0.
        center_covariances (bool, optional): Use centered covariances to compute the VAMP score. Defaults to True.

    """
    cov_X, cov_Y, cov_XY = (
        covariance(X, center=center_covariances),
        covariance(Y, center=center_covariances),
        covariance(X, Y, center=center_covariances),
    )
    R_X = log_fro_metric_deformation_loss(cov_X)
    R_Y = log_fro_metric_deformation_loss(cov_Y)
    if relaxed:
        S = (torch.linalg.matrix_norm(cov_XY, ord="fro") ** 2) / (
            (
                torch.linalg.matrix_norm(cov_X, ord=2)
                * torch.linalg.matrix_norm(cov_Y, ord=2)
            )
        )
    else:
        M_X = torch.linalg.lstsq(cov_X, cov_XY).solution
        M_Y = torch.linalg.lstsq(cov_Y, cov_XY.T).solution
        S = torch.trace(M_X @ M_Y)
    return S - 0.5 * metric_deformation * (R_X + R_Y)
//  [python]
def fit(self, train_kernel, train_targets, train_logits=None):
        self.num_labels = train_kernel.size(0)
        self.kernel_dtype = train_kernel.dtype

        kernel = self.get_regularized_kernel(train_kernel)
        train_targets_coords = self.get_target_coords(train_targets)

        if train_logits is not None and self.args.f0_scaling > 0:
            train_targets_coords -= train_logits / self.args.f0_scaling

        self.kernel_solution, self.residuals, self.rank, _ = torch.solve(train_targets_coords.t(), kernel)
//  [python]
def fit(self, train_kernel, train_targets, train_logits=None):
        self.num_labels = train_kernel.size(0)
        self.kernel_dtype = train_kernel.dtype

        kernel = self.get_regularized_kernel(train_kernel)
        train_targets_coords = self.get_target_coords(train_targets)

        if train_logits is not None and self.args.f0_scaling > 0:
            train_targets_coords -= train_logits / self.args.f0_scaling

        self.kernel_solution, self.residuals, self.rank, _ = torch.linalg.lstsq(kernel, train_targets_coords.t())
//  [python]
def apply_constraint_linear(
    org_params: torch.Tensor,
    param_map: ParamMap,
    key_point_map: torch.Tensor,
    constraints: torch.Tensor,
) -> torch.Tensor:
    """
    org_params: (n-dim) or (m, n-dim)
    constraints: (n-constraints, n-dim)
    key_point_map: (n-kp, n-dim)

    returns (1 or m, n-dim)
    """
    device = org_params.device

    if org_params.dim() == 1:
        org_params = org_params.unsqueeze(0)  # (m, n-dim)

    ns_basis = torch.tensor(scipy.linalg.null_space(constraints.cpu().numpy()), device=device).T  # (n-ns, n-dim)

    kp_from_ns = key_point_map @ ns_basis.T  # (n-kp, n-ns)

    org_kp = org_params @ key_point_map.T  # (m, n-kp)

    scale_factor = _compute_scale_factor(param_map, org_params).unsqueeze(-1)  # (m, n-kp, 1)

    # minimize (kp_from_ns) @ P - org_kp
    driver = "gels" if scale_factor.is_cuda else "gelsd"
    sol = torch.solve(
        org_kp.unsqueeze(-1) * scale_factor, kp_from_ns * scale_factor
    ).solution.squeeze(
        -1
    )  # (m, n-ns)
    return sol @ ns_basis
//  [python]
def apply_constraint_linear(
    org_params: torch.Tensor,
    param_map: ParamMap,
    key_point_map: torch.Tensor,
    constraints: torch.Tensor,
) -> torch.Tensor:
    """
    org_params: (n-dim) or (m, n-dim)
    constraints: (n-constraints, n-dim)
    key_point_map: (n-kp, n-dim)

    returns (1 or m, n-dim)
    """
    device = org_params.device

    if org_params.dim() == 1:
        org_params = org_params.unsqueeze(0)  # (m, n-dim)

    ns_basis = torch.tensor(scipy.linalg.null_space(constraints.cpu().numpy()), device=device).T  # (n-ns, n-dim)

    kp_from_ns = key_point_map @ ns_basis.T  # (n-kp, n-ns)

    org_kp = org_params @ key_point_map.T  # (m, n-kp)

    scale_factor = _compute_scale_factor(param_map, org_params).unsqueeze(-1)  # (m, n-kp, 1)

    # minimize (kp_from_ns) @ P - org_kp
    driver = "gels" if scale_factor.is_cuda else "gelsd"
    sol = torch.linalg.lstsq(
        kp_from_ns * scale_factor, org_kp.unsqueeze(-1) * scale_factor, driver=driver
    ).solution.squeeze(
        -1
    )  # (m, n-ns)
    return sol @ ns_basis
//  [python]
def f(x, y, lamb: float = 0):
    """f.
    Ridge polynomial regression. Referred from h16.py

    :param x:
    :param y:
    :param lamb:
    :type lamb: float
    """
    X = torch.vander(x, 3 + 1)
    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
    loss = F.mse_loss(y_hat, y) * len(y)
    return w, loss
//  [python]
def f(x, y, lamb: float = 0):
    """f.
    Ridge polynomial regression. Referred from h16.py

    :param x:
    :param y:
    :param lamb:
    :type lamb: float
    """
    X = torch.vander(x, 3 + 1)
    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
    loss = F.mse_loss(y_hat, y) * len(y)
    return w, loss
//  [python]
def one_step_forward(self, z, return_rec=False, return_K=False):
        B, input_len, E = z.shape
        assert input_len > 1, 'snapshots number should be larger than 1'
        x, y = z[:, :-1], z[:, 1:]

        # solve linear system
        self.K = torch.solve(torch.bmm(x.transpose(1, 2), y), torch.bmm(x.transpose(1, 2), x)).solution # B E E
        if torch.isnan(self.K).any():
            print('Encounter K with nan, replace K by identity matrix')
            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)

        z_pred = torch.bmm(z[:, -1:], self.K)
        if return_rec:
            z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1)
            return z_rec, z_pred

        return z_pred
//  [python]
def one_step_forward(self, z, return_rec=False, return_K=False):
        B, input_len, E = z.shape
        assert input_len > 1, 'snapshots number should be larger than 1'
        x, y = z[:, :-1], z[:, 1:]

        # solve linear system
        self.K = torch.linalg.lstsq(x, y).solution # B E E
        if torch.isnan(self.K).any():
            print('Encounter K with nan, replace K by identity matrix')
            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)

        z_pred = torch.bmm(z[:, -1:], self.K)
        if return_rec:
            z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1)
            return z_rec, z_pred

        return z_pred
//  [python]
def __call__(self, xs, ys, inds=None):
        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 = []

        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])

        return torch.stack(preds, dim=1)
//  [python]
def __call__(self, xs, ys, inds=None):
        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 = []

        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])

        return torch.stack(preds, dim=1)
//  [python]
def fit_AE(self, x):     
        if self.pos_encoding:
            x = torch.add(x, self.encoding) #simpler case, adds pos encoding
        # temp = torch.add(temp, self.pos_encoding)
        # temp = torch.mm(self._alpha, torch.cat((x, self.pos_encoding))) # new case, concatenate the encoding

        # H = self._activation(torch.add(temp, self.bias))
        H = self._activation(torch.mm(self._alpha, x))
        # self._beta = torch.mm(torch.mm(x, H.t()), torch.linalg.inv(torch.mm(H, H.t()) + self._lambda * self._eye))
        
        # self._beta = torch.mm(x, torch.pinverse(H))
        self._beta = torch.solve(x.t(), H.t()).solution
        return self._beta
//  [python]
def fit_AE(self, x):     
        if self.pos_encoding:
            x = torch.add(x, self.encoding) #simpler case, adds pos encoding
        # temp = torch.add(temp, self.pos_encoding)
        # temp = torch.mm(self._alpha, torch.cat((x, self.pos_encoding))) # new case, concatenate the encoding

        # H = self._activation(torch.add(temp, self.bias))
        H = self._activation(torch.mm(self._alpha, x))
        # self._beta = torch.mm(torch.mm(x, H.t()), torch.linalg.inv(torch.mm(H, H.t()) + self._lambda * self._eye))
        
        # self._beta = torch.mm(x, torch.pinverse(H))
        self._beta = torch.linalg.lstsq(H.t(), x.t()).solution
        return self._beta
//  [python]
def pointmap2fmap(p2p, evecs_x, evecs_y):
    """
    Convert a point-to-point map to functional map

    Args:
        p2p (np.ndarray): point-to-point map (shape x -> shape y). [Vx]
        evecs_x (np.ndarray): eigenvectors of shape x. [Vx, K]
        evecs_y (np.ndarray): eigenvectors of shape y. [Vy, K]
    Returns:
        C21 (np.ndarray): functional map (shape y -> shape x). [K, K]
    """
    C21 = torch.solve(evecs_y[p2p, :].T @ evecs_x, evecs_x.T @ evecs_x).solution
    return C21
//  [python]
def pointmap2fmap(p2p, evecs_x, evecs_y):
    """
    Convert a point-to-point map to functional map

    Args:
        p2p (np.ndarray): point-to-point map (shape x -> shape y). [Vx]
        evecs_x (np.ndarray): eigenvectors of shape x. [Vx, K]
        evecs_y (np.ndarray): eigenvectors of shape y. [Vy, K]
    Returns:
        C21 (np.ndarray): functional map (shape y -> shape x). [K, K]
    """
    C21 = torch.linalg.lstsq(evecs_x, evecs_y[p2p, :]).solution
    return C21
//  [python]
def forward(self, z, pred_len=1):
        # z:       B L E, koopman invariance space representation
        # z_rec:   B L E, reconstructed representation
        # z_pred:  B S E, forecasting representation
        B, input_len, E = z.shape
        assert input_len > 1, 'snapshots number should be larger than 1'
        x, y = z[:, :-1], z[:, 1:]

        # solve linear system
        self.K = torch.solve(torch.bmm(x.transpose(1, 2), y), torch.bmm(x.transpose(1, 2), x)).solution # B E E

        if torch.isnan(self.K).any():
            print('Encounter K with nan, replace K by identity matrix')
            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)

        z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1) # B L E
        
        if pred_len <= input_len:
            self.K_step = torch.linalg.matrix_power(self.K, pred_len)
            if torch.isnan(self.K_step).any():
                print('Encounter multistep K with nan, replace it by identity matrix')
                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
            z_pred = torch.bmm(z[:, -pred_len:, :], self.K_step)
        else:
            self.K_step = torch.linalg.matrix_power(self.K, input_len)
            if torch.isnan(self.K_step).any():
                print('Encounter multistep K with nan, replace it by identity matrix')
                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
            temp_z_pred, all_pred = z, []
            for _ in range(math.ceil(pred_len / input_len)):
                temp_z_pred = torch.bmm(temp_z_pred, self.K_step)
                all_pred.append(temp_z_pred)
            z_pred = torch.cat(all_pred, dim=1)[:, :pred_len, :]

        return z_rec, z_pred
//  [python]
def forward(self, z, pred_len=1):
        # z:       B L E, koopman invariance space representation
        # z_rec:   B L E, reconstructed representation
        # z_pred:  B S E, forecasting representation
        B, input_len, E = z.shape
        assert input_len > 1, 'snapshots number should be larger than 1'
        x, y = z[:, :-1], z[:, 1:]

        # solve linear system
        self.K = torch.linalg.lstsq(x, y).solution # B E E

        if torch.isnan(self.K).any():
            print('Encounter K with nan, replace K by identity matrix')
            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)

        z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1) # B L E
        
        if pred_len <= input_len:
            self.K_step = torch.linalg.matrix_power(self.K, pred_len)
            if torch.isnan(self.K_step).any():
                print('Encounter multistep K with nan, replace it by identity matrix')
                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
            z_pred = torch.bmm(z[:, -pred_len:, :], self.K_step)
        else:
            self.K_step = torch.linalg.matrix_power(self.K, input_len)
            if torch.isnan(self.K_step).any():
                print('Encounter multistep K with nan, replace it by identity matrix')
                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
            temp_z_pred, all_pred = z, []
            for _ in range(math.ceil(pred_len / input_len)):
                temp_z_pred = torch.bmm(temp_z_pred, self.K_step)
                all_pred.append(temp_z_pred)
            z_pred = torch.cat(all_pred, dim=1)[:, :pred_len, :]

        return z_rec, z_pred
//  [python]
def test_linalg_lstsq_batch_broadcasting(self, device, dtype):
        from torch.testing._internal.common_utils import random_well_conditioned_matrix

        def check_correctness(a, b):
            sol = torch.linalg.lstsq(a, b).solution
            sol2 = a.pinverse() @ b
            self.assertEqual(sol, sol2, rtol=1e-5, atol=1e-5)

        ms = [2 ** i for i in range(5)]
        batches = [(), (0,), (2,), (2, 2), (2, 2, 2)]
        # the case when a single matrix is batch-broadcasted over the rhs
        for m, batch in itertools.product(ms, batches):
            a = random_well_conditioned_matrix(m, m, dtype=dtype, device=device).view(*([1] * len(batch)), m, m)
            b = torch.rand(*(batch + (m, m)), dtype=dtype, device=device)
            check_correctness(a, b)

        # cases with broadcastable shapes
        for m in ms:
            a = random_well_conditioned_matrix(1, 3, 1, 3, m, m, dtype=dtype, device=device)
            b = torch.rand(3, 1, 3, 1, m, m // 2, dtype=dtype, device=device)
            check_correctness(a, b)

            # rhs are vectors, not matrices in this test
            b = torch.rand(3, 1, 3, 1, m, dtype=dtype, device=device)
            # unsqueeze for b because `check_correctness` checks against
            # a.pinverse() @ b, which requires b to be a matrix
            check_correctness(a, b.unsqueeze(-1))

            a = random_well_conditioned_matrix(3, 1, 3, 1, m, m, dtype=dtype, device=device)
            b = torch.rand(1, 3, 1, 3, m, m // 2, dtype=dtype, device=device)
            check_correctness(a, b)

            # rhs are vectors, not matrices in this test
            b = torch.rand(1, 3, 1, 3, m, dtype=dtype, device=device)
            check_correctness(a, b.unsqueeze(-1))
//  [python]
def least_squares_predictor(X, y):
    return torch.lstsq(y, X)[0]
//  [python]
def least_squares_predictor(X, y):
    return torch.linalg.lstsq(X, y).solution
//  [python]
def XA_b(A, B):
    x = torch.lstsq(B, A)[0]
    res = (A @ x - B)
    return x.T, res
//  [python]
def XA_b(A, B):
    x = torch.linalg.lstsq(A, B).solution
    res = (A @ x - B)
    return x.T, res
//  [python]
def fit_token(self, features, mask_low):
        """
        Finds token that when dot-producted with features minimizes MSE with low 
        resolution masks.

        Args:
            features (Nx256x64x64)
            mask (Nx1x64x64) - Should be logits type
        """
        with torch.no_grad():
            N = features.shape[0]
            assert N == mask_low.shape[0]
            A = features.permute(0, 2, 3, 1).reshape(N * 64 * 64, 256)
            B = mask_low.permute(0, 2, 3, 1).reshape(N * 64 * 64, 1)
            X = torch.solve(torch.matmul(A.t(),B), torch.matmul(A.t(),A)).solution.reshape(1, 256, 1, 1)
        return X.detach()
//  [python]
def fit_token(self, features, mask_low):
        """
        Finds token that when dot-producted with features minimizes MSE with low 
        resolution masks.

        Args:
            features (Nx256x64x64)
            mask (Nx1x64x64) - Should be logits type
        """
        with torch.no_grad():
            N = features.shape[0]
            assert N == mask_low.shape[0]
            A = features.permute(0, 2, 3, 1).reshape(N * 64 * 64, 256)
            B = mask_low.permute(0, 2, 3, 1).reshape(N * 64 * 64, 1)
            X = torch.linalg.lstsq(A, B).solution.reshape(1, 256, 1, 1)
        return X.detach()
//  [python]
def meanEffects(self):
        """The mean effects of each mutation against each phenotype, returned as a (p x D) tensor
        """

        X, y = self[: len(self)][:2]
        sol, _, _, _ = torch.lstsq(X, y)
        sol = sol.T

        return sol[: self.p, :]
//  [python]
def meanEffects(self):
        """The mean effects of each mutation against each phenotype, returned as a (p x D) tensor
        """

        X, y = self[: len(self)][:2]
        lstsq_ret = torch.linalg.lstsq(y, X)
        sol = lstsq_ret.solution

        return sol[: self.p, :]
//  [python]
def fit(self, episodes):
        # sequence_length * batch_size x feature_size
        featmat = self._feature(episodes).view(-1, self.feature_size)
        # sequence_length * batch_size x 1
        returns = episodes.returns.view(-1, 1)

        # Remove blank (all-zero) episodes that only exist because episode lengths vary
        flat_mask = episodes.mask.flatten()
        flat_mask_nnz = torch.nonzero(flat_mask)
        featmat = featmat[flat_mask_nnz].view(-1, self.feature_size)
        returns = returns[flat_mask_nnz].view(-1, 1)

        reg_coeff = self._reg_coeff
        XT_y = torch.matmul(featmat.t(), returns)
        XT_X = torch.matmul(featmat.t(), featmat)
        for _ in range(5):
            try:
                coeffs, _ = torch.lstsq(XT_y, XT_X + reg_coeff * self._eye)
                if (
                    torch.isnan(coeffs).any()
                    or torch.isinf(coeffs).any()
                ):
                    raise RuntimeError

                break
            except RuntimeError:
                reg_coeff *= 10
        else:
            raise RuntimeError(
                "Unable to solve the normal equations in "
                "`LinearFeatureBaseline`. The matrix X^T*X (with X the design "
                "matrix) is not full-rank, regardless of the regularization "
                "(maximum regularization: {0}).".format(reg_coeff)
            )
        self.weight.copy_(coeffs.flatten())
//  [python]
def fit(self, episodes):
        # sequence_length * batch_size x feature_size
        featmat = self._feature(episodes).view(-1, self.feature_size)
        # sequence_length * batch_size x 1
        returns = episodes.returns.view(-1, 1)

        # Remove blank (all-zero) episodes that only exist because episode lengths vary
        flat_mask = episodes.mask.flatten()
        flat_mask_nnz = torch.nonzero(flat_mask)
        featmat = featmat[flat_mask_nnz].view(-1, self.feature_size)
        returns = returns[flat_mask_nnz].view(-1, 1)

        reg_coeff = self._reg_coeff
        XT_y = torch.matmul(featmat.t(), returns)
        XT_X = torch.matmul(featmat.t(), featmat)
        for _ in range(5):
            try:
                # coeffs, _ = torch.lstsq(XT_y, XT_X + reg_coeff * self._eye
                # torch 1.x
                coeffs = torch.linalg.lstsq(XT_y, XT_X + reg_coeff * self._eye)

                # An extra round of increasing regularization eliminated
                # inf or nan in the least-squares solution most of the time
                if (
                    torch.isnan(coeffs.solution).any()
                    or torch.isinf(coeffs.solution).any()
                ):
                    raise RuntimeError

                break
            except RuntimeError:
                reg_coeff *= 10
        else:
            raise RuntimeError(
                "Unable to solve the normal equations in "
                "`LinearFeatureBaseline`. The matrix X^T*X (with X the design "
                "matrix) is not full-rank, regardless of the regularization "
                "(maximum regularization: {0}).".format(reg_coeff)
            )
        self.weight.copy_(coeffs.solution.flatten())
//  [python]
def _get_perspective_coeffs(
        startpoints: List[List[int]], endpoints: List[List[int]]
) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.solve(b_matrix.unsqueeze(1), a_matrix).solution

    output: List[float] = res.tolist()
    return output