jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def rq(
    torch: Any,
    tensor: Tensor,
    pivot_axis: int,
    non_negative_diagonal: bool = False
) -> Tuple[Tensor, Tensor]:
  """Computes the RQ decomposition of a tensor.

  The RQ decomposition is performed by treating the tensor as a matrix,
  with an effective left (row) index resulting from combining the axes
  `tensor.shape[:pivot_axis]` and an effective right (column) index
  resulting from combining the axes `tensor.shape[pivot_axis:]`.

  For example, if `tensor` had a shape (2, 3, 4, 5) and `pivot_axis` was 2,
  then `r` would have shape (2, 3, 6), and `q` would have shape (6, 4, 5).

  The output consists of two tensors `R, Q` such that:
  
    R[i1,...,iN, j] * Q[j, k1,...,kM] == tensor[i1,...,iN, k1,...,kM]
  
  `R` is a lower triangular matrix, `Q` is an orthonormal matrix
  Note that the output ordering matches numpy.linalg.svd rather than tf.svd.

  Args:
    tf: The tensorflow module.
    tensor: A tensor to be decomposed.
    pivot_axis: Where to split the tensor's axes before flattening into a
     p matrix.

  Returns:
    R: Left tensor factor.
    Q: Right tensor factor.
  """

  left_dims = tensor.shape[:pivot_axis]
  right_dims = tensor.shape[pivot_axis:]
  tensor = torch.reshape(tensor, [np.prod(left_dims), np.prod(right_dims)])
  #torch has currently no support for complex dtypes
  q, r = torch.linalg.qr(torch.transpose(tensor, 0, 1))
  if non_negative_diagonal:
    phases = torch.sign(torch.diagonal(r))
    q = q * phases
    r = phases[:, None] * r
  r, q = torch.transpose(r, 0, 1), torch.transpose(q, 0,
                                                   1)  #M=r*q at this point
  center_dim = r.shape[1]
  r = torch.reshape(r, list(left_dims) + [center_dim])
  q = torch.reshape(q, [center_dim] + list(right_dims))
  return r, q
//  [python]
def gauss_newton_step(betas,rho,L):
    M = betas.shape[0]
    A,b = computeJacobian(betas,rho,L)

    ata = torch.bmm(A.permute(0,2,1),A)
    q,r = torch.qr(A)
    qtb = torch.bmm(q.permute(0,2,1),b.unsqueeze(-1))
    r_inv = torch.inverse(r)

    rinv_qtb = torch.bmm(r_inv,qtb)
    next_betas = betas.unsqueeze(-1) + rinv_qtb

    error = torch.bmm(b.view((M,1,6)),b.view((M,6,1)))
    return next_betas.squeeze(-1), error
//  [python]
def gauss_newton_step(betas,rho,L):
    M = betas.shape[0]
    A,b = computeJacobian(betas,rho,L)

    ata = torch.bmm(A.permute(0,2,1),A)
    q,r = torch.linalg.qr(A)
    qtb = torch.bmm(q.permute(0,2,1),b.unsqueeze(-1))
    r_inv = torch.inverse(r)

    rinv_qtb = torch.bmm(r_inv,qtb)
    next_betas = betas.unsqueeze(-1) + rinv_qtb

    error = torch.bmm(b.view((M,1,6)),b.view((M,6,1)))
    return next_betas.squeeze(-1), error
//  [python]
def orthogonal(tensor, gain=1):
    """Fills the input Tensor or Variable with a (semi) orthogonal matrix, as described in "Exact solutions to the
    nonlinear dynamics of learning in deep linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional torch.Tensor or autograd.Variable, where n >= 2
        gain: optional scaling factor

    Examples:
        >>> w = torch.Tensor(3, 5)
        >>> nn.init.orthogonal(w)
    """
    if isinstance(tensor, Variable):
        orthogonal(tensor.data, gain=gain)
        return tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)
    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)
    # Pad zeros to Q (if rows smaller than cols)
    if rows < cols:
        padding = torch.zeros(rows, cols - rows)
        if q.is_cuda:
            q = torch.cat([q, padding.cuda()], 1)
        else:
            q = torch.cat([q, padding], 1)

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal(tensor, gain=1):
    """Fills the input Tensor or Variable with a (semi) orthogonal matrix, as described in "Exact solutions to the
    nonlinear dynamics of learning in deep linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional torch.Tensor or autograd.Variable, where n >= 2
        gain: optional scaling factor

    Examples:
        >>> w = torch.Tensor(3, 5)
        >>> nn.init.orthogonal(w)
    """
    if isinstance(tensor, Variable):
        orthogonal(tensor.data, gain=gain)
        return tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)
    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)
    # Pad zeros to Q (if rows smaller than cols)
    if rows < cols:
        padding = torch.zeros(rows, cols - rows)
        if q.is_cuda:
            q = torch.cat([q, padding.cuda()], 1)
        else:
            q = torch.cat([q, padding], 1)

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def __call__(self, raw_feat_s: Tensor, raw_feat_q: Tensor, **kwargs):
        all_features = torch.cat([raw_feat_s, raw_feat_q], 0)
        all_features = torch.qr(all_features.t()).R
        all_features = all_features.t()
        return all_features[: raw_feat_s.size(0)], all_features[raw_feat_s.size(0) :]
//  [python]
def __call__(self, raw_feat_s: Tensor, raw_feat_q: Tensor, **kwargs):
        all_features = torch.cat([raw_feat_s, raw_feat_q], 0)
        all_features = torch.linalg.qr(all_features.t()).R
        all_features = all_features.t()
        return all_features[: raw_feat_s.size(0)], all_features[raw_feat_s.size(0) :]
//  [python]
def __init__(self, in_channel):
        super().__init__()
        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel):
        super().__init__()
        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert(n_split % 2 == 0)
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian

        w_init = torch.qr(torch.FloatTensor(
            self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert(n_split % 2 == 0)
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian

        w_init = torch.linalg.qr(torch.FloatTensor(
            self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def init_orthogonal(tensor, gain=1):
    r"""
    Taken from a future torch version
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def init_orthogonal(tensor, gain=1):
    r"""
    Taken from a future torch version
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def process_gaussian(A: torch.Tensor, scale: float, inverted: bool):
    """
    Convert a Gaussian x = scale*A z, with z ~ N(0, I), into a linear autoregressive model
    Calling this function with (A, scale=s) is equivalent to calling it with (s*A, scale=1),
    except with possibly better numerical stability.
    """
    assert isinstance(A, torch.Tensor) and len(A.shape) == 2 and A.shape[0] == A.shape[1]
    dim = A.shape[0]
    if inverted:
        # Here, the goal is to convert x = A^{-1} z into a linear AR model.
        # If R (upper triangular) comes from the QR decomposition A, then
        # R^T R = A^T A, so A^{-1} A^{-1}^T = R^{-1} R^{-1}^T,
        # which implies that R^{-1} z has the same distribution as A^{-1} z.
        # So, if x solves the equation R x = z, then x will have the distribution we desire.
        # Because R is triangular, x can be obtained via back substitution, which is equivalent to sampling from
        # a linear AR model driven by Gaussian noise z. The coefficients of the linear functions that
        # define the conditionals of this AR model can be read off in the rows of R; the ordering of the
        # variables goes backwards because R is upper triangular.
        _, R = torch.qr(A)
        stds = (1.0 / scale) / R.diag().abs()  # standard deviations of conditionals
        R /= -R.diag()[:, None]  # coefficients of linear functions defining the means
        mean_coefs = R
    else:
        # Here, the goal is to convert x = A z into a linear AR model
        L = torch.qr(A.t())[1].t()  # Cholesky decomposition of AA^T
        stds = L.diag().abs() * scale  # standard deviations of conditionals
        Linv, _ = torch.triangular_solve(torch.eye(dim, dtype=L.dtype, device=L.device), L, upper=False)  # invert L
        Linv *= -L.diag()[:, None]
        mean_coefs = Linv
    mean_coefs[range(dim), range(dim)] = 0  # set diagonal to zero; AR conditionals don't depend on current timestep
    return mean_coefs, stds
//  [python]
def process_gaussian(A: torch.Tensor, scale: float, inverted: bool):
    """
    Convert a Gaussian x = scale*A z, with z ~ N(0, I), into a linear autoregressive model
    Calling this function with (A, scale=s) is equivalent to calling it with (s*A, scale=1),
    except with possibly better numerical stability.
    """
    assert isinstance(A, torch.Tensor) and len(A.shape) == 2 and A.shape[0] == A.shape[1]
    dim = A.shape[0]
    if inverted:
        # Here, the goal is to convert x = A^{-1} z into a linear AR model.
        # If R (upper triangular) comes from the QR decomposition A, then
        # R^T R = A^T A, so A^{-1} A^{-1}^T = R^{-1} R^{-1}^T,
        # which implies that R^{-1} z has the same distribution as A^{-1} z.
        # So, if x solves the equation R x = z, then x will have the distribution we desire.
        # Because R is triangular, x can be obtained via back substitution, which is equivalent to sampling from
        # a linear AR model driven by Gaussian noise z. The coefficients of the linear functions that
        # define the conditionals of this AR model can be read off in the rows of R; the ordering of the
        # variables goes backwards because R is upper triangular.
        _, R = torch.linalg.qr(A)
        stds = (1.0 / scale) / R.diag().abs()  # standard deviations of conditionals
        R /= -R.diag()[:, None]  # coefficients of linear functions defining the means
        mean_coefs = R
    else:
        # Here, the goal is to convert x = A z into a linear AR model
        L = torch.linalg.qr(A.t())[1].t()  # Cholesky decomposition of AA^T
        stds = L.diag().abs() * scale  # standard deviations of conditionals
        Linv, _ = torch.triangular_solve(torch.eye(dim, dtype=L.dtype, device=L.device), L, upper=False)  # invert L
        Linv *= -L.diag()[:, None]
        mean_coefs = Linv
    mean_coefs[range(dim), range(dim)] = 0  # set diagonal to zero; AR conditionals don't depend on current timestep
    return mean_coefs, stds
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian
        w_init = torch.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian
        w_init = torch.linalg.qr(torch.FloatTensor(self.n_split, self.n_split).normal_())[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def orthonormal_init(param: nn.Parameter, n_blocks: int):
        size0, size1 = param.size()
        size0 //= n_blocks
        size_min = min(size0, size1)
        init_values = []
        for _ in range(n_blocks):
            m1 = torch.randn(size0, size0, dtype=param.dtype)
            m2 = torch.randn(size1, size1, dtype=param.dtype)
            q1, r1 = torch.qr(m1)
            q2, r2 = torch.qr(m2)
            q1 *= torch.sign(torch.diag(r1))
            q2 *= torch.sign(torch.diag(r2))
            value = torch.mm(q1[:, :size_min], q2[:size_min, :])
            init_values.append(value)
        param.data = torch.cat(init_values, dim=0)
//  [python]
def orthonormal_init(param: nn.Parameter, n_blocks: int):
        size0, size1 = param.size()
        size0 //= n_blocks
        size_min = min(size0, size1)
        init_values = []
        for _ in range(n_blocks):
            m1 = torch.randn(size0, size0, dtype=param.dtype)
            m2 = torch.randn(size1, size1, dtype=param.dtype)
            q1, r1 = torch.linalg.qr(m1)
            q2, r2 = torch.linalg.qr(m2)
            q1 *= torch.sign(torch.diag(r1))
            q2 *= torch.sign(torch.diag(r2))
            value = torch.mm(q1[:, :size_min], q2[:size_min, :])
            init_values.append(value)
        param.data = torch.cat(init_values, dim=0)
//  [python]
def test_qr(self, device):
        def run_test(tensor_dims, some):
            A = torch.randn(*tensor_dims, device=device)
            Q, R = torch.qr(A, some=some)

            # Check0: Q[-2:] = (m, n_columns), R[-2:] = (n_columns, n)
            m, n = tensor_dims[-2:]
            n_columns = m if (not some) and m > n else min(m, n)
            self.assertEqual(Q.size(-2), m)
            self.assertEqual(R.size(-1), n)
            self.assertEqual(Q.size(-1), n_columns)

            # Check1: A = QR
            self.assertEqual(A, torch.matmul(Q, R))

            # Check2: A = QR (with out)
            Q_out, R_out = torch.Tensor().to(device), torch.Tensor().to(device)
            torch.qr(A, some=some, out=(Q_out, R_out))
            self.assertEqual(A, torch.matmul(Q_out, R_out))

            # Check3: Q == Q_out, R == R_out
            self.assertEqual(Q, Q_out)
            self.assertEqual(R, R_out)

            # Check4: Q^{T}Q = I, triu(R) = R
            self.assertEqual(torch.matmul(Q.transpose(-2, -1), Q),
                             torch.eye(n_columns, device=device).expand(Q.shape[:-2] + (n_columns, n_columns)))
            self.assertEqual(R.triu(), R)

        tensor_dims_list = [(3, 5), (5, 5), (5, 3),  # Single matrix
                            (7, 3, 5), (7, 5, 5), (7, 5, 3),  # 3-dim Tensors
                            (7, 5, 3, 5), (7, 5, 5, 5), (7, 5, 5, 3)]  # 4-dim Tensors
        for tensor_dims, some in product(tensor_dims_list, [True, False]):
            run_test(tensor_dims, some)
//  [python]
def test_qr(self, device):
        def run_test(tensor_dims, some):
            A = torch.randn(*tensor_dims, device=device)
            Q, R = torch.linalg.qr(A, mode='reduced' if not some else 'complete')

            # Check0: Q[-2:] = (m, n_columns), R[-2:] = (n_columns, n)
            m, n = tensor_dims[-2:]
            n_columns = m if (not some) and m > n else min(m, n)
            self.assertEqual(Q.size(-2), m)
            self.assertEqual(R.size(-1), n)
            self.assertEqual(Q.size(-1), n_columns)

            # Check1: A = QR
            self.assertEqual(A, torch.matmul(Q, R))

            # Check2: A = QR (with out)
            Q_out, R_out = torch.Tensor().to(device), torch.Tensor().to(device)
            torch.linalg.qr(A, mode='reduced' if not some else 'complete', out=(Q_out, R_out))
            self.assertEqual(A, torch.matmul(Q_out, R_out))

            # Check3: Q == Q_out, R == R_out
            self.assertEqual(Q, Q_out)
            self.assertEqual(R, R_out)

            # Check4: Q^{T}Q = I, triu(R) = R
            self.assertEqual(torch.matmul(Q.transpose(-2, -1), Q),
                             torch.eye(n_columns, device=device).expand(Q.shape[:-2] + (n_columns, n_columns)))
            self.assertEqual(R.triu(), R)

        tensor_dims_list = [(3, 5), (5, 5), (5, 3),  # Single matrix
                            (7, 3, 5), (7, 5, 5), (7, 5, 3),  # 3-dim Tensors
                            (7, 5, 3, 5), (7, 5, 5, 5), (7, 5, 5, 3)]  # 4-dim Tensors
        for tensor_dims, some in product(tensor_dims_list, [True, False]):
            run_test(tensor_dims, some)
//  [python]
def orthogonalize_weights(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    tensor = tensor.clone().detach() # @mst: avoid modifying the original tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    # flattened = tensor.new(rows, cols).normal_(0, 1)
    flattened = tensor.view(rows, cols) # @mst: do NOT reinit the tensor

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonalize_weights(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    tensor = tensor.clone().detach() # @mst: avoid modifying the original tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    # flattened = tensor.new(rows, cols).normal_(0, 1)
    flattened = tensor.view(rows, cols) # @mst: do NOT reinit the tensor

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_init(layers, mean=0.0, std=0.01):
    k = len(layers)
    ou_f = layers[0].out_features
    in_f = layers[0].in_features
    random = torch.randn((ou_f, in_f, k)) * std + mean
    q, r = torch.qr(random, some=True)

    for detector, init in zip(layers, q.permute(2, 0, 1)):
        detector.weight.data.copy_(init)
        nn.init.zeros_(detector.bias)
//  [python]
def orthogonal_init(layers, mean=0.0, std=0.01):
    k = len(layers)
    ou_f = layers[0].out_features
    in_f = layers[0].in_features
    random = torch.randn((ou_f, in_f, k)) * std + mean
    q, r = torch.linalg.qr(random, mode='reduced')

    for detector, init in zip(layers, q.permute(2, 0, 1)):
        detector.weight.data.copy_(init)
        nn.init.zeros_(detector.bias)
//  [python]
def lr_torch(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.qr(x2_flat)
    return (torch.norm(q2.T @ x1_flat)) ** 2 / (torch.norm(x1_flat)) ** 2
//  [python]
def lr_torch(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.linalg.qr(x2_flat)
    return (torch.norm(q2.T @ x1_flat)) ** 2 / (torch.norm(x1_flat)) ** 2
//  [python]
def _orthogonalize_batch_tt_cores_left_to_right(tt):
    """Orthogonalize TT-cores of a batch TT-object in the left to right order.
    Args:
        tt: TensorTrainBatch.
    Returns:
        TensorTrainBatch
    """
    # Left to right othogonalizatino.
    ndims = tt.ndims 
    raw_shape = tt.get_raw_shape()
    tt_ranks = tt.get_tt_ranks()
    next_rank = tt_ranks[0]
    batch_size = tt.batch_size 

    # Copy cores reference so we can change the cores. 
    tt_cores = list(tt.tt_cores)
    for core_idx in range(ndims - 1):
        curr_core = tt_cores[core_idx]
        # TT-ranks could have changed on the previous iteration, so `tt_ranks` can
        # be outdated for the current TT-rank, but should be valid for the next
        # TT-rank.
        curr_rank = next_rank 
        next_rank = tt_ranks[core_idx + 1]
        if tt.is_tt_matrix():
            curr_mode_left = raw_shape[0][core_idx]
            curr_mode_right = raw_shape[1][core_idx]
            curr_mode = curr_mode_left * curr_mode_right 
        else:
            curr_mode = raw_shape[0][core_idx]
        
        qr_shape = (batch_size, curr_rank * curr_mode, next_rank)
        curr_core = curr_core.reshape(qr_shape)
        curr_core, triang = torch.qr(curr_core)
        triang_shape = triang.shape 
        # The TT-rank could have changed: if qr_shape is e.g. 4 x 10, than q would
        # be of size 4 x 4 and r would be 4 x 10, which means that the next rank
        # should be changed to 4.
        next_rank = triang_shape[1]
        if tt.is_tt_matrix():
            new_core_shape = (batch_size, curr_rank, curr_mode_left, curr_mode_right, next_rank)
        else:
            new_core_shape = (batch_size, curr_rank, curr_mode, next_rank)
        
        tt_cores[core_idx] = curr_core.reshape(new_core_shape)
        next_core = tt_cores[core_idx + 1].reshape(batch_size, triang_shape[2], -1)
        tt_cores[core_idx + 1] = torch.mm(triang, next_core)
    
    if tt.is_tt_matrix():
        last_core_shape = (batch_size, next_rank, raw_shape[0][-1], raw_shape[1][-1], 1)
    else:
        last_core_shape = (batch_size, next_rank, raw_shape[0][-1], 1)
    tt_cores[-1] = tt_cores[-1].reshape(last_core_shape)
    # TODO: infer the tt_ranks.

    return TensorTrainBatch(tt_cores, tt.get_raw_shape(), batch_size=batch_size)
//  [python]
def _orthogonalize_batch_tt_cores_left_to_right(tt):
    """Orthogonalize TT-cores of a batch TT-object in the left to right order.
    Args:
        tt: TensorTrainBatch.
    Returns:
        TensorTrainBatch
    """
    # Left to right othogonalizatino.
    ndims = tt.ndims 
    raw_shape = tt.get_raw_shape()
    tt_ranks = tt.get_tt_ranks()
    next_rank = tt_ranks[0]
    batch_size = tt.batch_size 

    # Copy cores reference so we can change the cores. 
    tt_cores = list(tt.tt_cores)
    for core_idx in range(ndims - 1):
        curr_core = tt_cores[core_idx]
        # TT-ranks could have changed on the previous iteration, so `tt_ranks` can
        # be outdated for the current TT-rank, but should be valid for the next
        # TT-rank.
        curr_rank = next_rank 
        next_rank = tt_ranks[core_idx + 1]
        if tt.is_tt_matrix():
            curr_mode_left = raw_shape[0][core_idx]
            curr_mode_right = raw_shape[1][core_idx]
            curr_mode = curr_mode_left * curr_mode_right 
        else:
            curr_mode = raw_shape[0][core_idx]
        
        qr_shape = (batch_size, curr_rank * curr_mode, next_rank)
        curr_core = curr_core.reshape(qr_shape)
        curr_core, triang = torch.linalg.qr(curr_core)
        triang_shape = triang.shape 
        # The TT-rank could have changed: if qr_shape is e.g. 4 x 10, than q would
        # be of size 4 x 4 and r would be 4 x 10, which means that the next rank
        # should be changed to 4.
        next_rank = triang_shape[1]
        if tt.is_tt_matrix():
            new_core_shape = (batch_size, curr_rank, curr_mode_left, curr_mode_right, next_rank)
        else:
            new_core_shape = (batch_size, curr_rank, curr_mode, next_rank)
        
        tt_cores[core_idx] = curr_core.reshape(new_core_shape)
        next_core = tt_cores[core_idx + 1].reshape(batch_size, triang_shape[2], -1)
        tt_cores[core_idx + 1] = torch.mm(triang, next_core)
    
    if tt.is_tt_matrix():
        last_core_shape = (batch_size, next_rank, raw_shape[0][-1], raw_shape[1][-1], 1)
    else:
        last_core_shape = (batch_size, next_rank, raw_shape[0][-1], 1)
    tt_cores[-1] = tt_cores[-1].reshape(last_core_shape)
    # TODO: infer the tt_ranks.

    return TensorTrainBatch(tt_cores, tt.get_raw_shape(), batch_size=batch_size)
//  [python]
def orthogonal(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def _orthogonalize_tt_cores_right_to_left(tt):
    """Orthogonalize TT-cores of a TT-object in the right to left order.
    Args:
        tt: TenosorTrain or a TensorTrainBatch.
    Returns:
        The same type as the input `tt` (TenosorTrain or a TensorTrainBatch).
    """
    # Left to right orthogonalization.
    ndims = tt.ndims 
    raw_shape = tt.get_raw_shape()
    tt_ranks = tt.get_tt_ranks()
    prev_rank = tt_ranks[ndims]
    # Copy cores reference so we can change the cores. 
    tt_cores = list(tt.tt_cores)
    for core_idx in range(ndims - 1, 0, -1):
        curr_core = tt_cores[core_idx]
        # TT-ranks could have changed on the previous iteration, so `tt_ranks` can
        # be outdated for the current TT-rank, but should be valid for the next
        # TT-rank.
        curr_rank = prev_rank
        prev_rank = tt_ranks[core_idx]
        if tt.is_tt_matrix():
            curr_mode_left = raw_shape[0][core_idx]
            curr_mode_right = raw_shape[1][core_idx]
            curr_mode = curr_mode_left * curr_mode_right 
        else:
            curr_mode = raw_shape[0][core_idx]

        qr_shape = (prev_rank, curr_mode * curr_rank)
        curr_core = curr_core.reshape(qr_shape)
        curr_core, triang = torch.qr(curr_core.t())
        curr_core = curr_core.t()
        triang_shape = triang.shape 

        # The TT-rank could have changed: if qr_shape is e.g. 4 x 10, than q would
        # be of size 4 x 4 and r would be 4 x 10, which means that the next rank
        # should be changed to 4.
        prev_rank = triang_shape[1]
        if tt.is_tt_matrix():
            new_core_shape = (prev_rank, curr_mode_left, curr_mode_right, curr_rank)
        else:
            new_core_shape = (prev_rank, curr_mode, curr_rank)
        tt_cores[core_idx] = curr_core.reshape(new_core_shape)

        prev_core = tt_cores[core_idx-1].reshape(-1, triang_shape[0])
        tt_cores[core_idx - 1] = torch.mm(prev_core, triang)
    
    if tt.is_tt_matrix():
        first_core_shape = (1, raw_shape[0][0], raw_shape[1][0], prev_rank)
    else:
        first_core_shape = (1, raw_shape[0][0], prev_rank)
    tt_cores[0] = tt_cores[0].reshape(first_core_shape)
    # TODO: infer the tt_ranks
    return TensorTrain(tt_cores, tt.get_raw_shape())
//  [python]
def _orthogonalize_tt_cores_right_to_left(tt):
    """Orthogonalize TT-cores of a TT-object in the right to left order.
    Args:
        tt: TenosorTrain or a TensorTrainBatch.
    Returns:
        The same type as the input `tt` (TenosorTrain or a TensorTrainBatch).
    """
    # Left to right orthogonalization.
    ndims = tt.ndims 
    raw_shape = tt.get_raw_shape()
    tt_ranks = tt.get_tt_ranks()
    prev_rank = tt_ranks[ndims]
    # Copy cores reference so we can change the cores. 
    tt_cores = list(tt.tt_cores)
    for core_idx in range(ndims - 1, 0, -1):
        curr_core = tt_cores[core_idx]
        # TT-ranks could have changed on the previous iteration, so `tt_ranks` can
        # be outdated for the current TT-rank, but should be valid for the next
        # TT-rank.
        curr_rank = prev_rank
        prev_rank = tt_ranks[core_idx]
        if tt.is_tt_matrix():
            curr_mode_left = raw_shape[0][core_idx]
            curr_mode_right = raw_shape[1][core_idx]
            curr_mode = curr_mode_left * curr_mode_right 
        else:
            curr_mode = raw_shape[0][core_idx]

        qr_shape = (prev_rank, curr_mode * curr_rank)
        curr_core = curr_core.reshape(qr_shape)
        curr_core, triang = torch.linalg.qr(curr_core.t())
        curr_core = curr_core.t()
        triang_shape = triang.shape 

        # The TT-rank could have changed: if qr_shape is e.g. 4 x 10, than q would
        # be of size 4 x 4 and r would be 4 x 10, which means that the next rank
        # should be changed to 4.
        prev_rank = triang_shape[1]
        if tt.is_tt_matrix():
            new_core_shape = (prev_rank, curr_mode_left, curr_mode_right, curr_rank)
        else:
            new_core_shape = (prev_rank, curr_mode, curr_rank)
        tt_cores[core_idx] = curr_core.reshape(new_core_shape)

        prev_core = tt_cores[core_idx-1].reshape(-1, triang_shape[0])
        tt_cores[core_idx - 1] = torch.mm(prev_core, triang)
    
    if tt.is_tt_matrix():
        first_core_shape = (1, raw_shape[0][0], raw_shape[1][0], prev_rank)
    else:
        first_core_shape = (1, raw_shape[0][0], prev_rank)
    tt_cores[0] = tt_cores[0].reshape(first_core_shape)
    # TODO: infer the tt_ranks
    return TensorTrain(tt_cores, tt.get_raw_shape())
//  [python]
def __init__(
        self, channels, num_splits=4, no_jacobian=False, **kwargs
    ):  # pylint: disable=unused-argument
        super().__init__()
        assert num_splits % 2 == 0
        self.channels = channels
        self.num_splits = num_splits
        self.no_jacobian = no_jacobian
        self.weight_inv = None

        if LooseVersion(torch.__version__) < LooseVersion("1.9"):
            w_init = torch.qr(
                torch.FloatTensor(self.num_splits, self.num_splits).normal_()
            )[0]
        else:
            w_init = torch.linalg.qr(
                torch.FloatTensor(self.num_splits, self.num_splits).normal_(),
                "complete",
            )[0]

        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def __init__(
        self, channels, num_splits=4, no_jacobian=False, **kwargs
    ):  # pylint: disable=unused-argument
        super().__init__()
        assert num_splits % 2 == 0
        self.channels = channels
        self.num_splits = num_splits
        self.no_jacobian = no_jacobian
        self.weight_inv = None

        w_init = torch.linalg.qr(
                torch.FloatTensor(self.num_splits, self.num_splits).normal_(),
                "complete",
            )[0]

        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def forward(self, A):
        Q, R = torch.qr(A)
        self.save_for_backward(A, Q, R)
        return Q, R
//  [python]
def forward(self, A):
        Q, R = torch.linalg.qr(A)
        self.save_for_backward(A, Q, R)
        return Q, R
//  [python]
def QRreduction(datas):
    # ndatas = torch.linalg.qr(datas.permute(0,2,1), mode='reduced').R
    ndatas = torch.qr(datas.permute(0, 2, 1)).R
    ndatas = ndatas.permute(0, 2, 1)
    return ndatas
//  [python]
def QRreduction(datas):
    ndatas = torch.linalg.qr(datas.permute(0,2,1), mode='reduced').R
    ndatas = ndatas.permute(0, 2, 1)
    return ndatas
//  [python]
def _orthogonal(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    dt = flattened.dtype
    flattened = flattened.to(torch.float32)  # orthogonal init does not support fp16
    q, r = torch.qr(flattened)
    q, r = q.to(dtype=dt), r.to(dtype=dt)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def _orthogonal(tensor, gain=1):
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    dt = flattened.dtype
    flattened = flattened.to(torch.float32)  # orthogonal init does not support fp16
    q, r = torch.linalg.qr(flattened)
    q, r = q.to(dtype=dt), r.to(dtype=dt)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def randomly_rotate(X: torch.Tensor) -> torch.Tensor:
    """Randomly rotate d,n data matrix X"""

    d, n = X.shape
    z = torch.randn((d, d), dtype=X.dtype)
    q, r = torch.qr(z)
    d = torch.diag(r)
    ph = d / abs(d)
    rot_mat = q * ph
    return rot_mat @ X
//  [python]
def randomly_rotate(X: torch.Tensor) -> torch.Tensor:
    """Randomly rotate d,n data matrix X"""

    d, n = X.shape
    z = torch.randn((d, d), dtype=X.dtype)
    q, r = torch.linalg.qr(z)
    d = torch.diag(r)
    ph = d / abs(d)
    rot_mat = q * ph
    return rot_mat @ X
//  [python]
def __init__(self, c):
        super().__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, c):
        super().__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, dim, scramble=False, hetero=True, hidden=False, child=True,  ones=True, noise_identity=True):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2
        print ("ones" + str(ones))
        

        if ones:
            self.wxy = torch.eye(self.dim)
            if child:
                print ("child " + str(child))
                self.wyz = torch.eye(self.dim)
            else:
                self.wyz =torch.zeros(self.dim, self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            if child:
                self.wyz = torch.randn(self.dim, self.dim) / dim
            else:
                self.wyz =torch.zeros(self.dim, self.dim)
        if scramble:
            self.scramble, _ = torch.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            if noise_identity==0:
                print ("noise_identity " + str(noise_identity))
                self.whx = torch.randn(self.dim, self.dim) / dim 
                self.why = torch.randn(self.dim, self.dim) / dim 
                self.whz = torch.randn(self.dim, self.dim) / dim 
            else:
                if noise_identity==1:
                    print ("noise_identity " + str(noise_identity))
                    self.whx = torch.eye(self.dim, self.dim)
                    self.why = torch.eye(self.dim, self.dim)
                    self.whz = torch.eye(self.dim, self.dim)
                else:
                    if noise_identity==2:
                        print ("noise_identity " + str(noise_identity))
                        self.whx = torch.rand(self.dim, self.dim) / dim
                        self.why = torch.rand(self.dim, self.dim) / dim
                        self.whz = torch.rand(self.dim, self.dim) / dim                    

        else:
            self.whx = torch.eye(self.dim, self.dim) 
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def __init__(self, dim, scramble=False, hetero=True, hidden=False, child=True,  ones=True, noise_identity=True):
        self.hetero = hetero
        self.hidden = hidden
        self.dim = dim // 2
        print ("ones" + str(ones))
        

        if ones:
            self.wxy = torch.eye(self.dim)
            if child:
                print ("child " + str(child))
                self.wyz = torch.eye(self.dim)
            else:
                self.wyz =torch.zeros(self.dim, self.dim)
        else:
            self.wxy = torch.randn(self.dim, self.dim) / dim
            if child:
                self.wyz = torch.randn(self.dim, self.dim) / dim
            else:
                self.wyz =torch.zeros(self.dim, self.dim)
        if scramble:
            self.scramble, _ = torch.linalg.qr(torch.randn(dim, dim))
        else:
            self.scramble = torch.eye(dim)

        if hidden:
            if noise_identity==0:
                print ("noise_identity " + str(noise_identity))
                self.whx = torch.randn(self.dim, self.dim) / dim 
                self.why = torch.randn(self.dim, self.dim) / dim 
                self.whz = torch.randn(self.dim, self.dim) / dim 
            else:
                if noise_identity==1:
                    print ("noise_identity " + str(noise_identity))
                    self.whx = torch.eye(self.dim, self.dim)
                    self.why = torch.eye(self.dim, self.dim)
                    self.whz = torch.eye(self.dim, self.dim)
                else:
                    if noise_identity==2:
                        print ("noise_identity " + str(noise_identity))
                        self.whx = torch.rand(self.dim, self.dim) / dim
                        self.why = torch.rand(self.dim, self.dim) / dim
                        self.whz = torch.rand(self.dim, self.dim) / dim                    

        else:
            self.whx = torch.eye(self.dim, self.dim) 
            self.why = torch.zeros(self.dim, self.dim)
            self.whz = torch.zeros(self.dim, self.dim)
//  [python]
def qr(mat: T.FloatTensor) -> T.Tuple[T.FloatTensor]:
    """
    Compute the QR decomposition of a matrix.
    The QR decomposition factorizes a matrix A into a product
    A = QR of an orthonormal matrix Q and an upper triangular matrix R.
    Provides an orthonormalization of the columns of the matrix.

    Args:
        mat: A matrix.

    Returns:
        (Q, R): Tuple of tensors.

    """
    return torch.qr(mat)
//  [python]
def qr(mat: T.FloatTensor) -> T.Tuple[T.FloatTensor]:
    """
    Compute the QR decomposition of a matrix.
    The QR decomposition factorizes a matrix A into a product
    A = QR of an orthonormal matrix Q and an upper triangular matrix R.
    Provides an orthonormalization of the columns of the matrix.

    Args:
        mat: A matrix.

    Returns:
        (Q, R): Tuple of tensors.

    """
    return torch.linalg.qr(mat)
//  [python]
def _orthonormal_initializer(self, weights, gain=1.0):
		if len(weights.shape)>2:
			pdb.set_trace()
		device = weights.device
		dtype = weights.dtype
		# pdb.set_trace()
		shape0, shape1 = get_shape(weights)
		M1 = torch.randn(size=(shape0, shape0), dtype=dtype, device=device)
		M2 = torch.randn(size=(shape1, shape1), dtype=dtype, device=device)
		Q1, R1 = torch.qr(M1)	# let weights.shape= (s0,s1) and sm = min(s0, s1), then Q1:(s0,sm), R1:(sm,s1)
		Q2, R2 = torch.qr(M2)
		Q1 = Q1 * torch.sign(torch.diag(R1))
		Q2 = Q2 * torch.sign(torch.diag(R2))
		n_min = min(shape0, shape1)
		
		with torch.no_grad():
			q = torch.matmul(Q1[:, :n_min], Q2[:n_min, :])
			weights.view_as(q).copy_(q)
			weights.mul_(gain)
		return weights
//  [python]
def _orthonormal_initializer(self, weights, gain=1.0):
		if len(weights.shape)>2:
			pdb.set_trace()
		device = weights.device
		dtype = weights.dtype
		# pdb.set_trace()
		shape0, shape1 = get_shape(weights)
		M1 = torch.randn(size=(shape0, shape0), dtype=dtype, device=device)
		M2 = torch.randn(size=(shape1, shape1), dtype=dtype, device=device)
		Q1, R1 = torch.linalg.qr(M1)	# let weights.shape= (s0,s1) and sm = min(s0, s1), then Q1:(s0,sm), R1:(sm,s1)
		Q2, R2 = torch.linalg.qr(M2)
		Q1 = Q1 * torch.sign(torch.diag(R1))
		Q2 = Q2 * torch.sign(torch.diag(R2))
		n_min = min(shape0, shape1)
		
		with torch.no_grad():
			q = torch.matmul(Q1[:, :n_min], Q2[:n_min, :])
			weights.view_as(q).copy_(q)
			weights.mul_(gain)
		return weights
//  [python]
def _inv_matmul_preconditioner(self):
        """
        (Optional) define a preconditioner that can be used for linear systems, but not necessarily
        for log determinants. By default, this can call :meth:`~gpytorch.lazy.LazyTensor._preconditioner`.

        Returns:
            function: a function on x which performs P^{-1}(x)
        """
        base_precond, _, _ = self._preconditioner()

        if base_precond is not None:
            return base_precond
        elif beta_features.default_preconditioner.on():
            if hasattr(self, "_default_preconditioner_cache"):
                U, S, V = self._default_preconditioner_cache
            else:
                precond_basis_size = min(settings.max_preconditioner_size.value(), self.size(-1))
                random_basis = torch.randn(
                    self.batch_shape + torch.Size((self.size(-2), precond_basis_size)),
                    device=self.device,
                    dtype=self.dtype,
                )
                projected_mat = self._matmul(random_basis)
                proj_q = torch.qr(projected_mat)
                orthog_projected_mat = self._matmul(proj_q).transpose(-2, -1)
                U, S, V = torch.svd(orthog_projected_mat)
                U = proj_q.matmul(U)

                self._default_preconditioner_cache = (U, S, V)

            def preconditioner(v):
                res = V.transpose(-2, -1).matmul(v)
                res = (1 / S).unsqueeze(-1) * res
                res = U.matmul(res)
                return res

            return preconditioner
        else:
            return None
//  [python]
def _inv_matmul_preconditioner(self):
        """
        (Optional) define a preconditioner that can be used for linear systems, but not necessarily
        for log determinants. By default, this can call :meth:`~gpytorch.lazy.LazyTensor._preconditioner`.

        Returns:
            function: a function on x which performs P^{-1}(x)
        """
        base_precond, _, _ = self._preconditioner()

        if base_precond is not None:
            return base_precond
        elif beta_features.default_preconditioner.on():
            if hasattr(self, "_default_preconditioner_cache"):
                U, S, V = self._default_preconditioner_cache
            else:
                precond_basis_size = min(settings.max_preconditioner_size.value(), self.size(-1))
                random_basis = torch.randn(
                    self.batch_shape + torch.Size((self.size(-2), precond_basis_size)),
                    device=self.device,
                    dtype=self.dtype,
                )
                projected_mat = self._matmul(random_basis)
                proj_q = torch.linalg.qr(projected_mat)
                orthog_projected_mat = self._matmul(proj_q[0]).transpose(-2, -1)
                U, S, V = torch.svd(orthog_projected_mat)
                U = proj_q[0].matmul(U)

                self._default_preconditioner_cache = (U, S, V)

            def preconditioner(v):
                res = V.transpose(-2, -1).matmul(v)
                res = (1 / S).unsqueeze(-1) * res
                res = U.matmul(res)
                return res

            return preconditioner
        else:
            return None