jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=10,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : zt
            - op : ['down', 'mid', 'up']
            - block_idx : op == down, up : [0,1,2,3], op == mid : [0]
            - pooling : ['pixel-sum', 'channel-sum', 'single-channel', 'multiple-channel']
        Returns
            - h : hidden feature
        '''
        # necessary variables
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).shape
        print('h_shape : ', h_shape)

        # if h_shape[1] > 512:
        #     chunk_size = 10
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1


        c_i, w_i, h_i = x.size(1), x.size(2), x.size(3)
        c_o, w_o, h_o = h_shape[1], h_shape[2], h_shape[3]

        a = th.tensor(0., device=x.device)

        # Algorithm 1
        vT = th.randn(c_i*w_i*h_i, pca_rank, device=x.device)
        q, r = th.linalg.qr(vT)
        vT = q
        v = vT.T
        v = v.view(-1, c_i, w_i, h_i)

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = th.cat(u, dim=0)
            u = u.to(x.device)

            # time_s = time.time()
            # g = lambda a : get_h(x + a*v)
            # u = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacfwd t ==', time_e - time_s)

            g = lambda x : einsum(u, get_h(x), 'b c w h, i c w h -> b')
            v_ = th.autograd.functional.jacobian(g, x)
            v_ = v_.view(-1, c_i*w_i*h_i)

            _, s, v = th.linalg.svd(v_, full_matrices=False)
            v = v.view(-1, c_i, w_i, h_i)
            u = u.view(-1, c_o, w_o, h_o)
            
            convergence = th.dist(v_prev, v.detach().cpu())
            print(f'power method : {i}-th step convergence : ', convergence)
            if th.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
                print('reach convergence threshold : ', convergence)
                break

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        u, s, vT = u.view(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.view(-1, c_i*w_i*h_i).detach()
        return u, s, vT
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=10,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : zt
            - op : ['down', 'mid', 'up']
            - block_idx : op == down, up : [0,1,2,3], op == mid : [0]
            - pooling : ['pixel-sum', 'channel-sum', 'single-channel', 'multiple-channel']
        Returns
            - h : hidden feature
        '''
        # necessary variables
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).shape
        print('h_shape : ', h_shape)

        # if h_shape[1] > 512:
        #     chunk_size = 10
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1


        c_i, w_i, h_i = x.size(1), x.size(2), x.size(3)
        c_o, w_o, h_o = h_shape[1], h_shape[2], h_shape[3]

        a = th.tensor(0., device=x.device)

        # Algorithm 1
        vT = th.randn(c_i*w_i*h_i, pca_rank, device=x.device)
        vT, _ = th.linalg.qr(vT)
        v = vT.T
        v = v.view(-1, c_i, w_i, h_i)

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = th.cat(u, dim=0)
            u = u.to(x.device)

            # time_s = time.time()
            # g = lambda a : get_h(x + a*v)
            # u = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacfwd t ==', time_e - time_s)

            g = lambda x : einsum(u, get_h(x), 'b c w h, i c w h -> b')
            v_ = th.autograd.functional.jacobian(g, x)
            v_ = v_.view(-1, c_i*w_i*h_i)

            _, s, v = th.linalg.svd(v_, full_matrices=False)
            v = v.view(-1, c_i, w_i, h_i)
            u = u.view(-1, c_o, w_o, h_o)
            
            convergence = th.dist(v_prev, v.detach().cpu())
            print(f'power method : {i}-th step convergence : ', convergence)
            if th.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
                print('reach convergence threshold : ', convergence)
                break

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        u, s, vT = u.view(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.view(-1, c_i*w_i*h_i).detach()
        return u, s, vT
//  [python]
def inject_perp_feature(
    trainset,
    epsilon,
    V,
    batch_size=128,
    num_workers=2,
    shuffle=True,
):
    num_classes = V.shape[0]

    x = torch.from_numpy(trainset.data.transpose([0, 3, 1, 2])).type(torch.float) / 255.0
    y = torch.tensor(trainset.targets, dtype=torch.long)

    Q, _ = torch.qr(V.reshape([num_classes, -1]).T)

    P = Q.matmul(Q.T)
    x_proj = P.matmul(x.view([-1, 3072]).T)

    x_poison = x - x_proj.T.view([-1, 3, 32, 32])

    for t in range(num_classes):
        x_poison[y == t] += epsilon * V[t]

    poisonset = torch.utils.data.TensorDataset(x_poison, y)
    poisonloader = torch.utils.data.DataLoader(
        poisonset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True
    )

    return poisonloader
//  [python]
def inject_perp_feature(
    trainset,
    epsilon,
    V,
    batch_size=128,
    num_workers=2,
    shuffle=True,
):
    num_classes = V.shape[0]

    x = torch.from_numpy(trainset.data.transpose([0, 3, 1, 2])).type(torch.float) / 255.0
    y = torch.tensor(trainset.targets, dtype=torch.long)

    Q, _ = torch.linalg.qr(V.reshape([num_classes, -1]).T)

    P = Q.matmul(Q.T)
    x_proj = P.matmul(x.view([-1, 3072]).T)

    x_poison = x - x_proj.T.view([-1, 3, 32, 32])

    for t in range(num_classes):
        x_poison[y == t] += epsilon * V[t]

    poisonset = torch.utils.data.TensorDataset(x_poison, y)
    poisonloader = torch.utils.data.DataLoader(
        poisonset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True
    )

    return poisonloader
//  [python]
def __init__(self, d: int, x_min: float, x_max: float) -> None:
        super().__init__()
        self.x_min = x_min
        self.x_max = x_max
        M = torch.randn(d, d)
        Q, _ = torch.qr(M)
        self.register_buffer('Q', Q)
//  [python]
def __init__(self, d: int, x_min: float, x_max: float) -> None:
        super().__init__()
        self.x_min = x_min
        self.x_max = x_max
        M = torch.randn(d, d)
        Q, _ = torch.linalg.qr(M)
        self.register_buffer('Q', Q)
//  [python]
def orthogonal_(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:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    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_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(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:
        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK)
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    if tensor.numel() == 0:
        # no-op
        return tensor
    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_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def inject_reverse_perp_feature(
    trainset,
    epsilon,
    V,
    batch_size=128,
    num_workers=2,
    shuffle=True,
):
    num_classes = V.shape[0]

    x = torch.from_numpy(trainset.data.transpose([0, 3, 1, 2])).type(torch.float) / 255.0
    y = torch.tensor(trainset.targets, dtype=torch.long)

    Q, _ = torch.qr(V.reshape([num_classes, -1]).T)

    P = Q.matmul(Q.T)
    x_proj = P.matmul(x.view([-1, 3072]).T)

    x_poison = x - x_proj.T.view([-1, 3, 32, 32])

    for t in range(num_classes):
        x_poison[y == t] -= epsilon * V[t]

    poisonset = torch.utils.data.TensorDataset(x_poison, y)
    poisonloader = torch.utils.data.DataLoader(
        poisonset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True
    )

    return poisonloader
//  [python]
def inject_reverse_perp_feature(
    trainset,
    epsilon,
    V,
    batch_size=128,
    num_workers=2,
    shuffle=True,
):
    num_classes = V.shape[0]

    x = torch.from_numpy(trainset.data.transpose([0, 3, 1, 2])).type(torch.float) / 255.0
    y = torch.tensor(trainset.targets, dtype=torch.long)

    Q, _ = torch.linalg.qr(V.reshape([num_classes, -1]).T)

    P = Q.matmul(Q.T)
    x_proj = P.matmul(x.view([-1, 3072]).T)

    x_poison = x - x_proj.T.view([-1, 3, 32, 32])

    for t in range(num_classes):
        x_poison[y == t] -= epsilon * V[t]

    poisonset = torch.utils.data.TensorDataset(x_poison, y)
    poisonloader = torch.utils.data.DataLoader(
        poisonset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=True
    )

    return poisonloader
//  [python]
def xdiag(A, m: int):
    n, n1 = A.shape
    assert n == n1
    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Omega = (2 * torch.randint(2, size=(n, m), device=A.device) - 1).float()
    Y = linops.operator_matrix_product(A, Omega)
    Q, R = torch.qr(Y)

    Z = linops.operator_matrix_product(A.T, Q)
    T = Z.T @ Omega
    S = normalize_columns(torch.linalg.inv(R).T)
    dQZ = diag_of_AB(Q.T, Z.T)
    dQSSZ = diag_of_AB((Q @ S).T, (Z @ S).T)
    dOmegaQT = diag_of_AB(Omega.T, (Q @ T).T)
    dOmegaY = diag_of_AB(Omega.T, Y.T)
    dOmegaQSST = diag_of_AB(Omega.T, 
                            (Q @ S @ torch.diag(diag_of_AB(S, T))).T)

    return dQZ + (-dQSSZ + dOmegaY -dOmegaQT +dOmegaQSST) / m
//  [python]
def xdiag(A, m: int):
    n, n1 = A.shape
    assert n == n1
    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Omega = (2 * torch.randint(2, size=(n, m), device=A.device) - 1).float()
    Y = linops.operator_matrix_product(A, Omega)
    Q, R = torch.linalg.qr(Y)

    Z = linops.operator_matrix_product(A.T, Q)
    T = Z.T @ Omega
    S = normalize_columns(torch.linalg.inv(R).T)
    dQZ = diag_of_AB(Q.T, Z.T)
    dQSSZ = diag_of_AB((Q @ S).T, (Z @ S).T)
    dOmegaQT = diag_of_AB(Omega.T, (Q @ T).T)
    dOmegaY = diag_of_AB(Omega.T, Y.T)
    dOmegaQSST = diag_of_AB(Omega.T, 
                            (Q @ S @ torch.diag(diag_of_AB(S, T))).T)

    return dQZ + (-dQSSZ + dOmegaY -dOmegaQT +dOmegaQSST) / m
//  [python]
def qr(x, mode="reduced"):
    x = convert_to_tensor(x)
    if mode not in {"reduced", "complete"}:
        raise ValueError(
            "`mode` argument value not supported. "
            "Expected one of {'reduced', 'complete'}. "
            f"Received: mode={mode}"
        )
    x = convert_to_tensor(x)
    return torch.qr(x, mode=mode)
//  [python]
def qr(x, mode="reduced"):
    x = convert_to_tensor(x)
    if mode not in {"reduced", "complete"}:
        raise ValueError(
            "`mode` argument value not supported. "
            "Expected one of {'reduced', 'complete'}. "
            f"Received: mode={mode}"
        )
    x = convert_to_tensor(x)
    return torch.linalg.qr(x, mode=mode)
//  [python]
def sample_v(x: torch.Tensor, hutchinson_samples: int):
    """
    Sample a random vector v of shape (*x.shape, hutchinson_samples)
    with scaled orthonormal columns.

    The reference data is used for shape, device and dtype.

    :param x: Reference data.
    :param hutchinson_samples: Number of Hutchinson samples to draw.
    :return:
    """
    batch_size, total_dim = x.shape[0], prod(x.shape[1:])
    if hutchinson_samples > total_dim:
        raise ValueError(f"Too many Hutchinson samples: got {hutchinson_samples}, expected <= {total_dim}")
    v = torch.randn(batch_size, total_dim, hutchinson_samples, device=x.device, dtype=x.dtype)
    q, r = torch.qr(v)
    q = q.reshape(*x.shape, hutchinson_samples)
    return q * sqrt(total_dim)
//  [python]
def sample_v(x: torch.Tensor, hutchinson_samples: int):
    """
    Sample a random vector v of shape (*x.shape, hutchinson_samples)
    with scaled orthonormal columns.

    The reference data is used for shape, device and dtype.

    :param x: Reference data.
    :param hutchinson_samples: Number of Hutchinson samples to draw.
    :return:
    """
    batch_size, total_dim = x.shape[0], prod(x.shape[1:])
    if hutchinson_samples > total_dim:
        raise ValueError(f"Too many Hutchinson samples: got {hutchinson_samples}, expected <= {total_dim}")
    v = torch.randn(batch_size, total_dim, hutchinson_samples, device=x.device, dtype=x.dtype)
    q = torch.linalg.qr(v).Q.reshape(*x.shape, hutchinson_samples)
    return q * sqrt(total_dim)
//  [python]
def apply_random_se3(coords_in, atom_mask=None, translation_scale=1.0):
    # unbatched. center on the mean of CA coords
    coords_mean = coords_in[:, 1:2].mean(-3, keepdim=True)
    coords_in -= coords_mean
    random_rot, _ = torch.qr(torch.randn(3, 3))
    coords_in = coords_in @ random_rot
    random_trans = torch.randn_like(coords_mean) * translation_scale
    coords_in += random_trans
    if atom_mask is not None:
        coords_in = coords_in * atom_mask[..., None]
    return coords_in
//  [python]
def apply_random_se3(coords_in, atom_mask=None, translation_scale=1.0):
    # unbatched. center on the mean of CA coords
    coords_mean = coords_in[:, 1:2].mean(-3, keepdim=True)
    coords_in -= coords_mean
    random_rot, _ = torch.linalg.qr(torch.randn(3, 3))
    coords_in = coords_in @ random_rot
    random_trans = torch.randn_like(coords_mean) * translation_scale
    coords_in += random_trans
    if atom_mask is not None:
        coords_in = coords_in * atom_mask[..., None]
    return coords_in
//  [python]
def __call__(self, x):
        M = torch.randn(3, 3)
        Q, __ = torch.qr(M)
        return x @ Q
//  [python]
def __call__(self, x):
        M = torch.randn(3, 3)
        Q, __ = torch.linalg.qr(M)
        return x @ Q
//  [python]
def orthogonal_matrix_chunk(cols, device = None):
    from distutils.version import LooseVersion
    TORCH_GE_1_8_0 = LooseVersion(torch.__version__) >= LooseVersion('1.8.0')
    unstructured_block = torch.randn((cols, cols), device = device)
    if TORCH_GE_1_8_0:
        q, r = torch.qr(unstructured_block.cpu(), some = True)
    else:
        q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def orthogonal_matrix_chunk(cols, device = None):
    from distutils.version import LooseVersion
    TORCH_GE_1_8_0 = LooseVersion(torch.__version__) >= LooseVersion('1.8.0')
    unstructured_block = torch.randn((cols, cols), device = device)
    if TORCH_GE_1_8_0:
        q, r = torch.linalg.qr(unstructured_block.cpu(), mode = 'reduced')
    else:
        q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def compute_regress(target, input, alpha=0.05):
    # Multiple linear regression using least squares
    # alpha 跟 R2 没有关系
    n, channel = input.shape

    if target.shape[0] != n:
        raise ValueError(f'InvalidData, {target.shape[0]} != {n}')

    if torch.isnan(input).int().sum():
        raise ValueError(f'NaN in input')
    '''
    Q: -0.0141
    R: -21.2489
    '''
    Q, R = torch.qr(input)  
    # perm = torch.argsort(torch.abs(torch.diag(R)))
    # Q, R = torch.qr(input[:, perm])

    # 少另外两个空条件
    # p = sum(abs(torch.diag(R)) > max(n, channel) * torch.ones(R[0]) * 1e-6)
    LS_coef = torch.inverse(R) @ (Q.T @ target)  # 0.8853, ...
    # LS_coef = LS_coef[perm]
    # RI = torch.reciprocal(torch.eye(int(p)) / R)
    # nu = max(0, n-p) # 5759931
    target_hat = input @ LS_coef  # 4096, 1
    r = target - target_hat
    norm_r = norm(r, axis=(0, 1))  # 0.0763

    # if nu != 0:
    #     rmse = norm_r / np.sqrt(nu)
    # tval = tinv((1-alpha/2), nu)

    SSE = norm_r * norm_r
    TSS = torch.pow(norm(target - torch.mean(target), axis=(0, 1)), 2)
    r2 = 1 - SSE / TSS
    # print("QR grad: ", r2.requires_grad)

    return r2
//  [python]
def compute_regress(target, input, alpha=0.05):
    # Multiple linear regression using least squares
    # alpha 跟 R2 没有关系
    n, channel = input.shape

    if target.shape[0] != n:
        raise ValueError(f'InvalidData, {target.shape[0]} != {n}')

    if torch.isnan(input).int().sum():
        raise ValueError(f'NaN in input')
    '''
    Q: -0.0141
    R: -21.2489
    '''
    #Q, R = torch.qr(input)  # 4096, 4096, 4096, 69
    Q, R = torch.linalg.qr(input, 'reduced' if True else 'complete')
    # perm = torch.argsort(torch.abs(torch.diag(R)))
    # Q, R = torch.qr(input[:, perm])

    # 少另外两个空条件
    # p = sum(abs(torch.diag(R)) > max(n, channel) * torch.ones(R[0]) * 1e-6)
    LS_coef = torch.inverse(R) @ (Q.T @ target)  # 0.8853, ...
    # LS_coef = LS_coef[perm]
    # RI = torch.reciprocal(torch.eye(int(p)) / R)
    # nu = max(0, n-p) # 5759931
    target_hat = input @ LS_coef  # 4096, 1
    r = target - target_hat
    norm_r = norm(r, axis=(0, 1))  # 0.0763

    # if nu != 0:
    #     rmse = norm_r / np.sqrt(nu)
    # tval = tinv((1-alpha/2), nu)

    SSE = norm_r * norm_r
    TSS = torch.pow(norm(target - torch.mean(target), axis=(0, 1)), 2)
    r2 = 1 - SSE / TSS
    # print("QR grad: ", r2.requires_grad)

    return r2
//  [python]
def test_equivariance(self):
        hidden_features = 3
        edge_features = 4
        num_layers = 4
        attention = False
        use_tanh = True
        eb = EquivariantBlock(
            hidden_features, edge_features, num_layers, attention, use_tanh
        )
        x = torch.randn(10, hidden_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 5 + [False] * 5)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.qr(torch.randn(3, 3))

        x1, pos1 = eb(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = eb(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [python]
def test_equivariance(self):
        hidden_features = 3
        edge_features = 4
        num_layers = 4
        attention = False
        use_tanh = True
        eb = EquivariantBlock(
            hidden_features, edge_features, num_layers, attention, use_tanh
        )
        x = torch.randn(10, hidden_features)
        pos = torch.randn(10, 3)
        edge_index = torch.tensor(
            [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [1, 2, 3, 4, 5, 6, 7, 8, 9, 0]]
        )
        mask = torch.tensor([True] * 5 + [False] * 5)
        edge_attr = torch.randn(10, edge_features)

        R, _ = torch.linalg.qr(torch.randn(3, 3))

        x1, pos1 = eb(x, pos, edge_index, mask, edge_attr)
        x2, pos2 = eb(x, pos @ R, edge_index, mask, edge_attr)
        self.assertTrue(torch.allclose(x1, x2))
        self.assertTrue(torch.allclose(pos1 @ R, pos2))
//  [python]
def samples_from_haar_distribution(d, bsize, device, dtype) -> torch.Tensor:
    """Random orthogonal matrices drawn from O(d) Haar distribution
    Adopted from scipy.stats.ortho_group, which implements the algorithm described in
    Mezzadri, How to generate random matrices from the classical compact groups (2006)
    """
    z = torch.randn(bsize, d, d, device=device, dtype=dtype)
    q, r = torch.qr(z)
    d = r.diagonal(offset=0, dim1=-2, dim2=-1)
    q *= (d/abs(d))[..., None, :]
    return q
//  [python]
def samples_from_haar_distribution(d, bsize, device, dtype) -> torch.Tensor:
    """Random orthogonal matrices drawn from O(d) Haar distribution
    Adopted from scipy.stats.ortho_group, which implements the algorithm described in
    Mezzadri, How to generate random matrices from the classical compact groups (2006)
    """
    z = torch.randn(bsize, d, d, device=device, dtype=dtype)
    q, r = torch.linalg.qr(z)
    # The last two dimensions are the rows and columns of R matrices.
    # Extract the diagonals. Note that this eliminates a dimension.
    d = r.diagonal(offset=0, dim1=-2, dim2=-1)
    # Add back a dimension for proper broadcasting: we're dividing
    # each row of each R matrix by the diagonal of the R matrix.
    q *= (d/abs(d))[..., None, :]  # to broadcast properly
    return q
//  [python]
def orthogonal_matrix_chunk(cols, device = None, dtype=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    q, r = torch.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t().to(dtype)
//  [python]
def orthogonal_matrix_chunk(cols, device = None, dtype=None):
    unstructured_block = torch.randn((cols, cols), device=device)
    q, r = torch.linalg.qr(unstructured_block.cpu(), mode='reduced')
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t().to(dtype)
//  [python]
def generate_ps(self, inp, N, level=None):
        # print(inp.shape) # (218, 178, 3)
        if self.preprocess is not None:
            transp, mean, std = self.preprocess
            inp = inp.transpose(*transp)
            inp = (inp - mean) / std
        # print(inp.shape) # (3, 218, 178)

        C, H, W = inp.shape
        h_use, w_use = int(H / self.factor), int(W / self.factor)

        ps = []
        for _ in range(N):
            p_signal = torch.zeros_like(inp)
            for c in range(C):
                rv = torch.randn((h_use, w_use))
                rv_ortho, _ = torch.qr(rv, some=True)
                p_signal[c, :h_use, :w_use] = rv_ortho
            p_img = RGB_signal_idct(p_signal.cpu().numpy())
            ps.append(p_img)
        ps = torch.from_numpy(np.stack(ps, axis=0), device=inp.device)

        if self.preprocess is not None:
            rev_transp = np.argsort(transp)
            ps = ps * std
            ps = ps.transpose(0, *(rev_transp + 1))
        # print(ps.shape) # (100, 218, 178, 3)
        return ps
//  [python]
def generate_ps(self, inp, N, level=None):
        # print(inp.shape) # (218, 178, 3)
        if self.preprocess is not None:
            transp, mean, std = self.preprocess
            inp = inp.transpose(*transp)
            inp = (inp - mean) / std
        # print(inp.shape) # (3, 218, 178)

        C, H, W = inp.shape
        h_use, w_use = int(H / self.factor), int(W / self.factor)

        ps = []
        for _ in range(N):
            p_signal = torch.zeros_like(inp)
            for c in range(C):
                rv = torch.randn((h_use, w_use))
                rv_ortho, _ = torch.linalg.qr(rv, mode="complete")
                p_signal[c, :h_use, :w_use] = rv_ortho
            p_img = RGB_signal_idct(p_signal.cpu().numpy())
            ps.append(p_img)
        ps = torch.from_numpy(np.stack(ps, axis=0), device=inp.device)

        if self.preprocess is not None:
            rev_transp = np.argsort(transp)
            ps = ps * std
            ps = ps.transpose(0, *(rev_transp + 1))
        # print(ps.shape) # (100, 218, 178, 3)
        return ps
//  [python]
def test_aten_fallback_must_fallback(self):
        class ModelWithAtenNotONNXOp(torch.nn.Module):
            def forward(self, x, y):
                abcd = x + y
                defg = torch.linalg.qr(abcd)
                return defg

        x = torch.rand(3, 4)
        y = torch.rand(3, 4)
        f = io.BytesIO()
        torch.onnx.export(
            ModelWithAtenNotONNXOp(),
            (x, y),
            f,
            do_constant_folding=False,
            operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,
            # support for linalg.qr was added in later op set versions.
            opset_version=9,
        )
        onnx_model = onnx.load(io.BytesIO(f.getvalue()))
        self.assertAtenOp(onnx_model, "linalg_qr")
//  [python]
def test_caffe2_aten_fallback_must_fallback(self):
        class ModelWithAtenNotONNXOp(torch.nn.Module):
            def forward(self, x, y):
                abcd = x + y
                defg = torch.qr(abcd)
                return defg

        # TODO: Refactor common_utils._decide_skip_caffe2 to support parametrize
        for operator_export_type in (
            OperatorExportTypes.ONNX_ATEN,
            OperatorExportTypes.ONNX_ATEN_FALLBACK,
        ):
            x = torch.rand(3, 4)
            y = torch.rand(3, 4)
            f = io.BytesIO()
            torch.onnx.export(
                ModelWithAtenNotONNXOp(),
                (x, y),
                f,
                do_constant_folding=False,
                operator_export_type=operator_export_type,
                # support for linalg.qr was added in later op set versions.
                opset_version=9,
            )
            onnx_model = onnx.load(io.BytesIO(f.getvalue()))
            self.assertAtenOp(onnx_model, "linalg_qr")
//  [python]
def test_caffe2_aten_fallback_must_fallback(self):
        class ModelWithAtenNotONNXOp(torch.nn.Module):
            def forward(self, x, y):
                abcd = x + y
                defg = torch.linalg.qr(abcd)
                return defg

        # TODO: Refactor common_utils._decide_skip_caffe2 to support parametrize
        for operator_export_type in (
            OperatorExportTypes.ONNX_ATEN,
            OperatorExportTypes.ONNX_ATEN_FALLBACK,
        ):
            x = torch.rand(3, 4)
            y = torch.rand(3, 4)
            f = io.BytesIO()
            torch.onnx.export(
                ModelWithAtenNotONNXOp(),
                (x, y),
                f,
                do_constant_folding=False,
                operator_export_type=operator_export_type,
                # support for linalg.qr was added in later op set versions.
                opset_version=9,
            )
            onnx_model = onnx.load(io.BytesIO(f.getvalue()))
            self.assertAtenOp(onnx_model, "linalg_qr")
//  [python]
def stable_qr(mat):
    """
    performs a QR decomposition on the batched matrix mat.
    We need to use these functions because of

    1. slow batched QR in pytorch (pytorch/pytorch#22573)
    2. possible singularity in R
    """
    if mat.shape[-1] <= stable_qr_cpu_threshold.value():
        # Dispatch to CPU so long as pytorch/pytorch#22573 is not fixed
        device = mat.device
        Q, R = torch.qr(mat.cpu())
        Q = Q.to(device)
        R = R.to(device)
    else:
        Q, R = torch.qr(mat)

    Rdiag = torch.diagonal(R, dim1=-2, dim2=-1)
    # if R is almost singular, add jitter
    zeroish = Rdiag.abs() < 1e-6
    if torch.any(zeroish):
        # can't use in-place operation here b/c it would mess up backward pass
        # haven't found a more elegant way to add a jitter diagonal yet...
        Rdiag_sign = torch.sign(Rdiag)
        # force zero diagonals to have jitter added to them.
        Rdiag_sign[Rdiag_sign == 0] = 1.0
        jitter_diag = 1e-6 * Rdiag_sign * zeroish.to(Rdiag)
        R = R + torch.diag_embed(jitter_diag)
    return Q, R
//  [python]
def stable_qr(mat):
    """
    performs a QR decomposition on the batched matrix mat.
    We need to use these functions because of

    1. slow batched QR in pytorch (pytorch/pytorch#22573)
    2. possible singularity in R
    """
    if mat.shape[-1] <= stable_qr_cpu_threshold.value():
        # Dispatch to CPU so long as pytorch/pytorch#22573 is not fixed
        device = mat.device
        Q, R = torch.linalg.qr(mat.cpu())
        Q = Q.to(device)
        R = R.to(device)
    else:
        Q, R = torch.linalg.qr(mat)

    Rdiag = torch.diagonal(R, dim1=-2, dim2=-1)
    # if R is almost singular, add jitter
    zeroish = Rdiag.abs() < 1e-6
    if torch.any(zeroish):
        # can't use in-place operation here b/c it would mess up backward pass
        # haven't found a more elegant way to add a jitter diagonal yet...
        Rdiag_sign = torch.sign(Rdiag)
        # force zero diagonals to have jitter added to them.
        Rdiag_sign[Rdiag_sign == 0] = 1.0
        jitter_diag = 1e-6 * Rdiag_sign * zeroish.to(Rdiag)
        R = R + torch.diag_embed(jitter_diag)
    return Q, R
//  [python]
def hutch_pp(solution: torch.Tensor, input_vec: torch.Tensor, m: int=102):
    """
    Algorithm taken from reference paper above on Hutch++.
    """
    assert m % 3 == 0
    if solution.numel() <= m:
        return exact_divergence(solution, input_vec)
    vjp = vjp_oracle_factory(solution, input_vec)
    k = m // 3
    S = 2.0 * torch.randint(0, 2, (*solution.shape, k), device=input_vec.device) - 1.0
    G = 2.0 * torch.randint(0, 2, (*solution.shape, k), device=input_vec.device) - 1.0
    BS = torch.empty_like(S)
    for i in range(k):
        BS[:, i] = vjp(S[:,i])
 
    Q, _ = torch.qr(BS)
    G_prime = G - Q @ (Q.T @ G)
    total = 0
    for i in range(Q.shape[1]):
        z = Q[:,i]
        total = total + (z * vjp(z)).sum()

    for i in range(k - 1):
        z = G_prime[:,i]
        total = total + (z * vjp(z)).sum() / k
    z = G_prime[:,-1]
    total = total + (z * vjp(z, False)).sum() / k

    return total
//  [python]
def hutch_pp(solution: torch.Tensor, input_vec: torch.Tensor, m: int=102):
    """
    Algorithm taken from reference paper above on Hutch++.
    """
    assert m % 3 == 0
    if solution.numel() <= m:
        return exact_divergence(solution, input_vec)
    vjp = vjp_oracle_factory(solution, input_vec)
    k = m // 3
    S = 2.0 * torch.randint(0, 2, (*solution.shape, k), device=input_vec.device) - 1.0
    G = 2.0 * torch.randint(0, 2, (*solution.shape, k), device=input_vec.device) - 1.0
    BS = torch.empty_like(S)
    for i in range(k):
        BS[:, i] = vjp(S[:,i])
 
    Q, _ = torch.linalg.qr(BS)
    G_prime = G - Q @ (Q.T @ G)
    total = 0
    for i in range(Q.shape[1]):
        z = Q[:,i]
        total = total + (z * vjp(z)).sum()

    for i in range(k - 1):
        z = G_prime[:,i]
        total = total + (z * vjp(z)).sum() / k
    z = G_prime[:,-1]
    total = total + (z * vjp(z, False)).sum() / k

    return total
//  [python]
def __init__(self, channels=3, permutation=None):
        super(GeneralizedChannelPermute, self).__init__()
        self.__delattr__("permutation")

        # Sample a random orthogonal matrix
        W, _ = torch.qr(torch.randn(channels, channels))

        # Construct the partially pivoted LU-form and the pivots
        LU, pivots = torch.lu_factor(W)

        # Convert the pivots into the permutation matrix
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != channels:
                raise ValueError(
                    'Keyword argument "permutation" expected to have {} elements but {} found.'.format(
                        channels, len(permutation)
                    )
                )
            P = torch.eye(channels, channels)[permutation.type(dtype=torch.int64)]

        # We register the permutation matrix so that the model can be serialized
        self.register_buffer("permutation", P)

        # NOTE: For this implementation I have chosen to store the parameters densely, rather than
        # storing L, U, and s separately
        self.LU = torch.nn.Parameter(LU)
//  [python]
def __init__(self, channels=3, permutation=None):
        super(GeneralizedChannelPermute, self).__init__()
        self.__delattr__("permutation")

        # Sample a random orthogonal matrix
        W, _ = torch.linalg.qr(torch.randn(channels, channels))

        # Construct the partially pivoted LU-form and the pivots
        LU, pivots = torch.linalg.lu_factor(W)

        # Convert the pivots into the permutation matrix
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != channels:
                raise ValueError(
                    'Keyword argument "permutation" expected to have {} elements but {} found.'.format(
                        channels, len(permutation)
                    )
                )
            P = torch.eye(channels, channels)[permutation.type(dtype=torch.int64)]

        # We register the permutation matrix so that the model can be serialized
        self.register_buffer("permutation", P)

        # NOTE: For this implementation I have chosen to store the parameters densely, rather than
        # storing L, U, and s separately
        self.LU = torch.nn.Parameter(LU)
//  [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)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        W = W.contiguous()
        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)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:, 0] = -1 * W[:, 0]
        W = W.view(c, c, 1)
        W = W.contiguous()
        self.conv.weight.data = W
//  [python]
def orthogonal_matrix_chunk_batched(bsz, cols, device=None):
    unstructured_block = torch.randn((bsz, cols, cols), device=device)
    q, r = torch.qr(unstructured_block)
    return q.transpose(2, 1)
//  [python]
def orthogonal_matrix_chunk_batched(bsz, cols, device=None):
    unstructured_block = torch.randn((bsz, cols, cols), device=device)
    q, r = torch.linalg.qr(unstructured_block, mode='reduced')
    return q.transpose(2, 1)
//  [python]
def lstq(A, F_u, F_v, lamda=0.01):
    requires_grad = F_u.requires_grad and F_u.grad_fn is not None # Requires grad and has tape attached
    if requires_grad and torch.__version__ >= '1.13.0' or not requires_grad and torch.__version__ >= '1.10.0':
        # 1.11.0 hopefully introduces grads to lstsq.
        theta_x = torch.linalg.lstsq(A, F_u).solution
        theta_y = torch.linalg.lstsq(A, F_v).solution
        return theta_x, theta_y

    # cols = A.shape[2]
    # assert all(cols == torch.linalg.matrix_rank(A))  # something better?
    try:
        Q, R = torch.qr(A)
        theta_x = torch.bmm(torch.bmm(torch.inverse(R), Q.transpose(1, 2)), F_u)
        theta_y = torch.bmm(torch.bmm(torch.inverse(R), Q.transpose(1, 2)), F_v)
    except:
        LOGGER.exception("Least Squares failed")
        raise
    return theta_x, theta_y
//  [python]
def lstq(A, F_u, F_v, lamda=0.01):
    requires_grad = F_u.requires_grad and F_u.grad_fn is not None # Requires grad and has tape attached
    if requires_grad and torch.__version__ >= '1.13.0' or not requires_grad and torch.__version__ >= '1.10.0':
        # 1.11.0 hopefully introduces grads to lstsq.
        theta_x = torch.linalg.lstsq(A, F_u).solution
        theta_y = torch.linalg.lstsq(A, F_v).solution
        return theta_x, theta_y

    # cols = A.shape[2]
    # assert all(cols == torch.linalg.matrix_rank(A))  # something better?
    try:
        Q, R = torch.linalg.qr(A)
        theta_x = torch.bmm(torch.bmm(torch.linalg.inv(R), Q.transpose(1, 2)), F_u)
        theta_y = torch.bmm(torch.bmm(torch.linalg.inv(R), Q.transpose(1, 2)), F_v)
    except:
        LOGGER.exception("Least Squares failed")
        raise
    return theta_x, theta_y
//  [python]
def M_to_G(self, x):
        g, r = torch.qr(x)
        r_diag = torch.diagonal(r, dim1=-2, dim2=-1)
        r_diag = torch.cat((r_diag, torch.ones((x.shape[0], self.n_m), dtype=dtype, device=device)), dim=1)
        g = g * r_diag[:, None]
        diff = 2*(torch.all(torch.isclose(g[:, :, :self.m], x), dim=-1).type(dtype)-0.5)
        g = g * diff[..., None]
        det_sign_g = torch.sign(torch.det(g))
        g[:, :, -1] *= det_sign_g[:, None]
        assert torch.allclose(x, g[:, :, :x.shape[-1]])
        assert torch.allclose(torch.det(g), torch.ones((g.shape[0],), dtype=dtype, device=device))
        return g
//  [python]
def M_to_G(self, x):
        g, r = torch.linalg.qr(x, mode='complete')
        r_diag = torch.diagonal(r, dim1=-2, dim2=-1)
        r_diag = torch.cat((r_diag, torch.ones((x.shape[0], self.n_m), dtype=dtype, device=device)), dim=1)
        g = g * r_diag[:, None]
        diff = 2*(torch.all(torch.isclose(g[:, :, :self.m], x), dim=-1).type(dtype)-0.5)
        g = g * diff[..., None]
        det_sign_g = torch.sign(torch.det(g))
        g[:, :, -1] *= det_sign_g[:, None]
        assert torch.allclose(x, g[:, :, :x.shape[-1]])
        assert torch.allclose(torch.det(g), torch.ones((g.shape[0],), dtype=dtype, device=device))
        return g
//  [python]
def sample_orthogonal_random_vectors(num_rows, num_cols, device=None):
        num_blocks = math.ceil(num_rows / num_cols)
        unstructured_block = torch.randn(num_blocks, num_cols, num_cols, device=device)
        norm = unstructured_block.norm(dim=-1).view(-1, 1)
        q, r = torch.qr(unstructured_block)

        # select num_rows vectors
        q = q.transpose(-1, -2).view(-1, num_cols)
        q = q[:num_rows, :]
        norm = norm[:num_rows, :]

        # multiply unit vectors by norm of random normal vectors
        # norm = torch.distributions.Chi2(q.new_tensor(num_cols)).sample((num_rows, 1)).sqrt()

        return q * norm
//  [python]
def sample_orthogonal_random_vectors(num_rows, num_cols, device=None):
        num_blocks = math.ceil(num_rows / num_cols)
        unstructured_block = torch.randn(num_blocks, num_cols, num_cols, device=device)
        norm = unstructured_block.norm(dim=-1).view(-1, 1)
        q, r = torch.linalg.qr(unstructured_block, mode='reduced')

        # select num_rows vectors
        q = q.transpose(-1, -2).view(-1, num_cols)
        q = q[:num_rows, :]
        norm = norm[:num_rows, :]

        # multiply unit vectors by norm of random normal vectors
        # norm = torch.distributions.Chi2(q.new_tensor(num_cols)).sample((num_rows, 1)).sqrt()

        return q * norm
//  [python]
def generate_observations(
    sample_size,
    mean_vector: torch.tensor,
    uv_scale: torch.tensor,
    mv_scale: torch.tensor,
) -> torch.tensor:
    """
    Generate two sequences with a correlation of `correlation`
    """
    if (
        mean_vector.shape != uv_scale.shape
        or mean_vector.shape[0] != mv_scale.shape[0]
        or mv_scale.shape[0] != mv_scale.shape[1]
    ):
        raise ValueError("Vectors must conform")

    torch.manual_seed(MANUAL_SEED)
    white_noise = torch.randn((sample_size, mean_vector.shape[0]))
    # Correct sample mean, sample variance, and orthogonality slightly so that the
    # sample statistics are what we want.
    white_noise = white_noise - torch.mean(white_noise, dim=0)
    q, r = torch.qr(white_noise)
    white_noise = q * sqrt(sample_size)

    # Resulting sample should have 1) zero mean; 2) unit variance; 3) orthogonal columns

    # Expand everyhing to be either (sample_size, n, 1) or (sample_size, n, n) to make
    # dimensions conform and to remove any ambiguity in the matrix multiply.

    white_noise_expanded = white_noise.unsqueeze(2).expand(
        (sample_size, mean_vector.shape[0], 1)
    )
    mean_vector = (
        mean_vector.unsqueeze(0).unsqueeze(2).expand(white_noise_expanded.shape)
    )
    uv_scale = uv_scale.unsqueeze(0).unsqueeze(2).expand(white_noise_expanded.shape)
    # Construct sample matrix and expand scale_matrix along sample dimension
    scale_matrix = mv_scale.unsqueeze(0).expand((sample_size,) + mv_scale.shape)

    # Multiply and drop the last dimension.
    random_observations = (
        uv_scale * (scale_matrix @ white_noise_expanded) + mean_vector
    ).squeeze(2)

    return random_observations, white_noise
//  [python]
def generate_observations(
    sample_size,
    mean_vector: torch.tensor,
    uv_scale: torch.tensor,
    mv_scale: torch.tensor,
) -> torch.tensor:
    """
    Generate two sequences with a correlation of `correlation`
    """
    if (
        mean_vector.shape != uv_scale.shape
        or mean_vector.shape[0] != mv_scale.shape[0]
        or mv_scale.shape[0] != mv_scale.shape[1]
    ):
        raise ValueError("Vectors must conform")

    torch.manual_seed(MANUAL_SEED)
    white_noise = torch.randn((sample_size, mean_vector.shape[0]))
    # Correct sample mean, sample variance, and orthogonality slightly so that the
    # sample statistics are what we want.
    white_noise = white_noise - torch.mean(white_noise, dim=0)
    white_noise = torch.linalg.qr(white_noise, mode="reduced")[0] * sqrt(sample_size)

    # Resulting sample should have 1) zero mean; 2) unit variance; 3) orthogonal columns

    # Expand everyhing to be either (sample_size, n, 1) or (sample_size, n, n) to make
    # dimensions conform and to remove any ambiguity in the matrix multiply.

    white_noise_expanded = white_noise.unsqueeze(2).expand(
        (sample_size, mean_vector.shape[0], 1)
    )
    mean_vector = (
        mean_vector.unsqueeze(0).unsqueeze(2).expand(white_noise_expanded.shape)
    )
    uv_scale = uv_scale.unsqueeze(0).unsqueeze(2).expand(white_noise_expanded.shape)
    # Construct sample matrix and expand scale_matrix along sample dimension
    scale_matrix = mv_scale.unsqueeze(0).expand((sample_size,) + mv_scale.shape)

    # Multiply and drop the last dimension.
    random_observations = (
        uv_scale * (scale_matrix @ white_noise_expanded) + mean_vector
    ).squeeze(2)

    return random_observations, white_noise
//  [python]
def qr(a: ArrayLike, mode="reduced"):
    a = _atleast_float_1(a)
    result = torch.qr(a, mode=mode)
    if mode == "r":
        # match NumPy
        result = result.R
    return result
//  [python]
def qr(a: ArrayLike, mode="reduced"):
    a = _atleast_float_1(a)
    result = torch.linalg.qr(a, mode=mode)
    if mode == "r":
        # match NumPy
        result = result.R
    return result
//  [python]
def test_orth_q(rng_seeds, torch_devices, heights_widths, entry_atol):
    """Test case for orthogonality of Q in QR decomposition.

    Tests that ``Q.T @ Q`` is close to identity.
    """
    for h, w in heights_widths:
        assert h >= w, "This test doesn't need/admit fat matrices!"
    #
    for seed in rng_seeds:
        for device in torch_devices:
            for dtype in (torch.float64, torch.float32):
                for h, w in heights_widths:
                    mat = gaussian_noise(
                        (h, w),
                        mean=0.0,
                        std=1.0,
                        seed=seed,
                        dtype=dtype,
                        device=device,
                    )
                    Q, R = torch.qr(mat)
                    I_residual = Q.T @ Q
                    I_residual[range(w), range(w)] -= 1
                    worst = I_residual.abs().max().item()
                    assert worst <= abs(entry_atol), "Q matrix not orthogonal?"
//  [python]
def test_orth_q(rng_seeds, torch_devices, heights_widths, entry_atol):
    """Test case for orthogonality of Q in QR decomposition.

    Tests that ``Q.T @ Q`` is close to identity.
    """
    for h, w in heights_widths:
        assert h >= w, "This test doesn't need/admit fat matrices!"
    #
    for seed in rng_seeds:
        for device in torch_devices:
            for dtype in (torch.float64, torch.float32):
                for h, w in heights_widths:
                    mat = gaussian_noise(
                        (h, w),
                        mean=0.0,
                        std=1.0,
                        seed=seed,
                        dtype=dtype,
                        device=device,
                    )
                    Q = torch.linalg.qr(mat)[0]
                    I_residual = Q.T @ Q
                    I_residual[range(w), range(w)] -= 1
                    worst = I_residual.abs().max().item()
                    assert worst <= abs(entry_atol), "Q matrix not orthogonal?"
//  [python]
def qr(self, a):
        return torch.qr(a)
//  [python]
def qr(self, a):
        return torch.linalg.qr(a)
//  [python]
def lstq(A, F_u, F_v, lamda=0.01):
    # cols = A.shape[2]
    # assert all(cols == torch.linalg.matrix_rank(A))  # something better?
    try:
        Q, R = torch.qr(A)
        theta_x = torch.bmm(torch.bmm(torch.inverse(R), Q.transpose(1, 2)), F_u)
        theta_y = torch.bmm(torch.bmm(torch.inverse(R), Q.transpose(1, 2)), F_v)
    except:
        LOGGER.exception("Least Squares failed")
        sys.exit(-1)
    return theta_x, theta_y