jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve_output(self, feature, target):
        ftf = feature.T @ feature
        fty = feature.T @ target

        try:
            output_w = torch.solve(fty, ftf)[0].to(self.opts.device)
        except:
            output_w = torch.mm(torch.pinverse(ftf), fty).to(self.opts.device)
            
        return output_w
//  [python]
def solve_output(self, feature, target):
        ftf = feature.T @ feature
        fty = feature.T @ target

        try:
            output_w = torch.linalg.lstsq(ftf, fty, driver = 'gelsd').solution.to(self.opts.device)
        except:
            output_w = torch.mm(torch.linalg.pinv(ftf), fty).to(self.opts.device)
            
        return output_w
//  [python]
def computetheta_ols(self, grid, flow, pred) :
        """
        Computes the parameter theta for this flow and weight map with vdist given
        Args :
            grid : (c, t, i, j, ft), c=2 and ft depends on the model
            flow : flow field (b, c, t, i, j)
            pred : weightning of different layers (b, l, t, i, j)
        Return :
            theta : ( b, l, ft) parameters for each layers and element of the batch
        """
        sc = ShapeCheck(pred.shape, 'b l t i j')
        sc.update(flow.shape, 'b c t i j')
        sc.update(grid.shape, 'c t i j ft')

        # Formulate preds to Weights
        pred = sc.repeat(pred, 'b l t i j -> (b l) (c t i j) 1')

        grid = sc.repeat(grid, 'c t i j ft -> (b l) (c t i j) ft')

        # Formulate Flow
        flow = sc.repeat(flow, 'b c t i j -> (b l) (c t i j) 1')

        try :
            wsq =  torch.sqrt(pred)
            Theta = torch.solve(torch.matmul(grid*wsq, (grid*wsq).T), grid*wsq, flow*wsq)[0]

            Theta = sc.repeat(Theta, '(b l) ft 1 -> b l ft')

        except Exception as e:
           Theta = torch.zeros(tuple(sc.get('b l ft').values()),
                   device=pred.device, requires_grad=True)
           print(f'Inversion Error : {e} batch')
           set_trace()
        if torch.isnan(Theta).sum() > 0 :
            print('Nan in Theta')
        return Theta.nan_to_num(0)
//  [python]
def computetheta_ols(self, grid, flow, pred) :
        """
        Computes the parameter theta for this flow and weight map with vdist given
        Args :
            grid : (c, t, i, j, ft), c=2 and ft depends on the model
            flow : flow field (b, c, t, i, j)
            pred : weightning of different layers (b, l, t, i, j)
        Return :
            theta : ( b, l, ft) parameters for each layers and element of the batch
        """
        sc = ShapeCheck(pred.shape, 'b l t i j')
        sc.update(flow.shape, 'b c t i j')
        sc.update(grid.shape, 'c t i j ft')

        # Formulate preds to Weights
        pred = sc.repeat(pred, 'b l t i j -> (b l) (c t i j) 1')

        grid = sc.repeat(grid, 'c t i j ft -> (b l) (c t i j) ft')

        # Formulate Flow
        flow = sc.repeat(flow, 'b c t i j -> (b l) (c t i j) 1')

        try :
            wsq =  torch.sqrt(pred)
            Theta = torch.linalg.lstsq(grid*wsq, flow*wsq).solution

            Theta = sc.repeat(Theta, '(b l) ft 1 -> b l ft')

        except Exception as e:
           Theta = torch.zeros(tuple(sc.get('b l ft').values()),
                   device=pred.device, requires_grad=True)
           print(f'Inversion Error : {e} batch')
           set_trace()
        if torch.isnan(Theta).sum() > 0 :
            print('Nan in Theta')
        return Theta.nan_to_num(0)
//  [python]
def forward(self, Z):
        # Z is in b * t x c x 1 x 1
        Zr = Z.squeeze().reshape(-1, self.n_frames, self.k_dim)

        if self.training and self.args.noise in ["latent"]:
            Zr = Zr + 0.003 * torch.rand(Zr.shape).to(Zr.device)

        # split
        X, Y = Zr[:, :-1], Zr[:, 1:]

        # solve linear system (broadcast)
        # Ct = torch.linalg.pinv(X.reshape(-1, self.k_dim)) @ Y.reshape(-1, self.k_dim)
        Ct = torch.solve(Y.reshape(-1, self.k_dim).T, X.reshape(-1, self.k_dim).T).solution.T

        # predict (broadcast)
        Y2 = X @ Ct
        Z2 = torch.cat((X[:, 0].unsqueeze(dim=1), Y2), dim=1)

        assert (torch.sum(torch.isnan(Y2)) == 0)

        return Z2.reshape(Z.shape), Ct
//  [python]
def forward(self, Z):
        # Z is in b * t x c x 1 x 1
        Zr = Z.squeeze().reshape(-1, self.n_frames, self.k_dim)

        if self.training and self.args.noise in ["latent"]:
            Zr = Zr + 0.003 * torch.rand(Zr.shape).to(Zr.device)

        # split
        X, Y = Zr[:, :-1], Zr[:, 1:]

        # solve linear system (broadcast)
        # Ct = torch.linalg.pinv(X.reshape(-1, self.k_dim)) @ Y.reshape(-1, self.k_dim)
        Ct = torch.linalg.lstsq(X.reshape(-1, self.k_dim), Y.reshape(-1, self.k_dim)).solution

        # predict (broadcast)
        Y2 = X @ Ct
        Z2 = torch.cat((X[:, 0].unsqueeze(dim=1), Y2), dim=1)

        assert (torch.sum(torch.isnan(Y2)) == 0)

        return Z2.reshape(Z.shape), Ct
//  [python]
def ray_estimation_loss(x, y, d_meas):  # for each ray
    # x as depth
    # y as sdf prediction
    # d_meas as measured depth

    # print(x.shape, y.shape, d_meas.shape)

    # regard each sample as a ray
    mat_A = torch.vstack((x, torch.ones_like(x))).transpose(0, 1)
    vec_b = y.view(-1, 1)

    # print(mat_A.shape, vec_b.shape)

    least_square_estimate = torch.lstsq(vec_b, mat_A).solution

    a = least_square_estimate[0]  # -> -1 (added in ekional loss term)
    b = least_square_estimate[1]

    # d_estimate = -b/a
    d_estimate = torch.clamp(-b / a, min=1.0, max=40.0)  # -> d

    # d_error = (d_estimate-d_meas)**2

    d_error = torch.abs(d_estimate - d_meas)

    # print(mat_A.shape, vec_b.shape, least_square_estimate.shape)
    # print(d_estimate.item(), d_meas.item(), d_error.item())

    return d_error
//  [python]
def ray_estimation_loss(x, y, d_meas):  # for each ray
    # x as depth
    # y as sdf prediction
    # d_meas as measured depth

    # print(x.shape, y.shape, d_meas.shape)

    # regard each sample as a ray
    mat_A = torch.vstack((x, torch.ones_like(x))).transpose(0, 1)
    vec_b = y.view(-1, 1)

    # print(mat_A.shape, vec_b.shape)

    least_square_estimate = torch.linalg.lstsq(mat_A, vec_b).solution

    a = least_square_estimate[0]  # -> -1 (added in ekional loss term)
    b = least_square_estimate[1]

    # d_estimate = -b/a
    d_estimate = torch.clamp(-b / a, min=1.0, max=40.0)  # -> d

    # d_error = (d_estimate-d_meas)**2

    d_error = torch.abs(d_estimate - d_meas)

    # print(mat_A.shape, vec_b.shape, least_square_estimate.shape)
    # print(d_estimate.item(), d_meas.item(), d_error.item())

    return d_error
//  [python]
def get_reward(self):
       
        probs = self.policy_net(self.last_state)
        m = Categorical(probs)      # 生成分布
        loss_entropy = -m.entropy() * 0.1
        with torch.no_grad():
            feat = self.extract_feat(self.student, self.last_state)
            if isinstance(feat, tuple):
                feat = feat[-1]
            # student_logit = self.task_head(feat)
        reward = - torch.solve(F.one_hot(self.last_target), feat).residuals
        # reward = - F.cross_entropy(student_logit, self.last_target, reduction='none').detach()
        loss_policy = -(m.log_prob(self.last_action) * reward).mean()

        return loss_policy, loss_entropy
//  [python]
def get_reward(self):
       
        probs = self.policy_net(self.last_state)
        m = Categorical(probs)      # 生成分布
        loss_entropy = -m.entropy() * 0.1
        with torch.no_grad():
            feat = self.extract_feat(self.student, self.last_state)
            if isinstance(feat, tuple):
                feat = feat[-1]
            # student_logit = self.task_head(feat)
        reward = - torch.linalg.lstsq(feat, F.one_hot(self.last_target)).residuals
        # reward = - F.cross_entropy(student_logit, self.last_target, reduction='none').detach()
        loss_policy = -(m.log_prob(self.last_action) * reward).mean()

        return loss_policy, loss_entropy
//  [python]
def _lstsq(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    if A.shape[0] != B.shape[0]:
        raise RuntimeError(f'A has different number of rows than B! A.shape = {A.shape}, B.shape = {B.shape}')
    if A.ndim != 2:
        raise RuntimeError('A is not a rank 2 tensor: has shape', A.shape)
    if B.ndim != 2:
        raise RuntimeError('B is not a rank 2 tensor: has shape', A.shape)

    # TODO more intelligence regarding choice of lstsq `driver` arg
    return torch.lstsq(A, B)[0]
//  [python]
def _lstsq(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    if A.shape[0] != B.shape[0]:
        raise RuntimeError(f'A has different number of rows than B! A.shape = {A.shape}, B.shape = {B.shape}')
    if A.ndim != 2:
        raise RuntimeError('A is not a rank 2 tensor: has shape', A.shape)
    if B.ndim != 2:
        raise RuntimeError('B is not a rank 2 tensor: has shape', A.shape)

    # TODO more intelligence regarding choice of lstsq `driver` arg
    return torch.linalg.lstsq(A, B).solution
//  [python]
def fit_bspline_to_traj(X, nland, s=None, include_resid=False):
    """Produce a B-spline from a trajectory with least-squares.

    Parameters
    ----------
    X : Tensor
        [neval,2] input trajectory
    nland : int
        number of landmarks (control points)
    s : Tensor
        (optional) [neval] time points for spline evaluation
    include_resid : bool
        whether to return the residuals of the least-squares problem

    Returns
    -------
    Y : Tensor
        [neval,2] output spline
    residuals : Tensor
        [2,] residuals of the least-squares problem (optional)
    """
    _check_input(X)
    neval = X.size(0)

    C = coefficient_mat(nland, neval, s=s, device=X.device)
    Y, residuals, _, _ = torch.lstsq(X, C)

    if include_resid:
        return Y, residuals

    return Y
//  [python]
def fit_bspline_to_traj(X, nland, s=None, include_resid=False):
    """Produce a B-spline from a trajectory with least-squares.

    Parameters
    ----------
    X : Tensor
        [neval,2] input trajectory
    nland : int
        number of landmarks (control points)
    s : Tensor
        (optional) [neval] time points for spline evaluation
    include_resid : bool
        whether to return the residuals of the least-squares problem

    Returns
    -------
    Y : Tensor
        [neval,2] output spline
    residuals : Tensor
        [2,] residuals of the least-squares problem (optional)
    """
    _check_input(X)
    neval = X.size(0)

    C = coefficient_mat(nland, neval, s=s, device=X.device)
    Y, residuals, _, _ = torch.linalg.lstsq(C, X, driver='gels')

    if include_resid:
        return Y, residuals

    return Y
//  [python]
def rectify_embeddings_by_FM(self):
        for scan_id in tqdm(range(self.start_index, self.end_index + 1)):
            scan_name = f"{self.subject_name}_{scan_id:05d}"

            fn_fitted_coarse = join(self.coarse_folder, f'optimized_template_{scan_name}.ply')
            fn_target_eigenfuncs = join(self.eigenfuncs_folder, f'{scan_name}_eigenfuncs.npz')


            try:
                print(f'[LOG] loading {scan_name}')
                template_fitted_verts, _ = igl.read_triangle_mesh(fn_fitted_coarse)
            except:
                print(f'[LOG] loading failed. skipping {scan_name}')
                continue

            template_fitted_verts = torch.from_numpy(template_fitted_verts).float().cuda()


            target_eigendata = np.load(fn_target_eigenfuncs)
            target_verts = torch.from_numpy(target_eigendata['verts']).float().cuda()
            target_eigenfuncs = torch.from_numpy(target_eigendata['eigenfuncs']).float().cuda()

            _, knn_idx_tem2tar, _ = knn_points(template_fitted_verts[None], target_verts[None])
            knn_idx_tem2tar = knn_idx_tem2tar[0, :, 0]
            eigenfunc_pullback_to_template = target_eigenfuncs[knn_idx_tem2tar]
            

            A = eigenfunc_pullback_to_template[:, 1:]
            B = self.template_eigenfuncs[:, 1:]
            FM = torch.solve(torch.matmul(A.t(),B),torch.matmul(A.t(),A)).solution

            target_eigen_corrected = target_eigenfuncs[:, 1:] @ FM

            # concatenate the constant back
            target_eigen_corrected = torch.cat([target_eigenfuncs[:, :1], target_eigen_corrected], -1)

            np.savez(join(self.eigenfuncs_retified_folder, scan_name + '_rectified.npz'),
                eigenfuncs=target_eigen_corrected.cpu().numpy(),
                eigenvals=target_eigendata['eigenvals'],
                verts=target_eigendata['verts'],
                faces=target_eigendata['faces'],
            )

            validate = True

            if validate:
                for i in range(target_eigen_corrected.shape[1]):
                    export_to_ply(join(self.eigenfuncs_retified_vis_folder, scan_name + f'_rectified_{i:03d}.ply'),
                                target_eigendata['verts'], target_eigendata['faces'],
                                target_eigen_corrected.cpu().numpy()[:, i])
//  [python]
def rectify_embeddings_by_FM(self):
        for scan_id in tqdm(range(self.start_index, self.end_index + 1)):
            scan_name = f"{self.subject_name}_{scan_id:05d}"

            fn_fitted_coarse = join(self.coarse_folder, f'optimized_template_{scan_name}.ply')
            fn_target_eigenfuncs = join(self.eigenfuncs_folder, f'{scan_name}_eigenfuncs.npz')


            try:
                print(f'[LOG] loading {scan_name}')
                template_fitted_verts, _ = igl.read_triangle_mesh(fn_fitted_coarse)
            except:
                print(f'[LOG] loading failed. skipping {scan_name}')
                continue

            template_fitted_verts = torch.from_numpy(template_fitted_verts).float().cuda()


            target_eigendata = np.load(fn_target_eigenfuncs)
            target_verts = torch.from_numpy(target_eigendata['verts']).float().cuda()
            target_eigenfuncs = torch.from_numpy(target_eigendata['eigenfuncs']).float().cuda()

            _, knn_idx_tem2tar, _ = knn_points(template_fitted_verts[None], target_verts[None])
            knn_idx_tem2tar = knn_idx_tem2tar[0, :, 0]
            eigenfunc_pullback_to_template = target_eigenfuncs[knn_idx_tem2tar]
            

            A = eigenfunc_pullback_to_template[:, 1:]
            B = self.template_eigenfuncs[:, 1:]
            FM = torch.linalg.lstsq(A, B).solution

            target_eigen_corrected = target_eigenfuncs[:, 1:] @ FM

            # concatenate the constant back
            target_eigen_corrected = torch.cat([target_eigenfuncs[:, :1], target_eigen_corrected], -1)

            np.savez(join(self.eigenfuncs_retified_folder, scan_name + '_rectified.npz'),
                eigenfuncs=target_eigen_corrected.cpu().numpy(),
                eigenvals=target_eigendata['eigenvals'],
                verts=target_eigendata['verts'],
                faces=target_eigendata['faces'],
            )

            validate = True

            if validate:
                for i in range(target_eigen_corrected.shape[1]):
                    export_to_ply(join(self.eigenfuncs_retified_vis_folder, scan_name + f'_rectified_{i:03d}.ply'),
                                target_eigendata['verts'], target_eigendata['faces'],
                                target_eigen_corrected.cpu().numpy()[:, i])
//  [python]
def test(self, nodes: Tuple[str], graph: BaseGraphInterface) -> TestResult:
        """
        Calculate the direct effect of each edge in the graph using multivariate regression.
        :param nodes: list of nodes
        :return: A TestResult with the action to take
        """
        effect_variable = graph.nodes[nodes[1]]
        cause_variable = graph.nodes[nodes[0]]

        if graph.undirected_edge_exists(cause_variable, effect_variable):
            # if the edge is not undirected, we do not need to calculate the direct effect
            return

        edge_data = graph.edge_value(cause_variable, effect_variable)

        all_parents_of_effect_variable = graph.parents_of_node(effect_variable)
        parents_of_effect_variable_without_cause_variable = [
            parent
            for parent in all_parents_of_effect_variable
            if parent != cause_variable
        ]

        effect_variable_parents_values = torch.stack(
            [cause_variable.values]
            + [
                parents_of_effect_variable_without_cause_variable[i].values
                for i in range(len(parents_of_effect_variable_without_cause_variable))
            ],
            dim=1,
        )
        # multivariate regression via tridiagonal reduction and SVD, see torch documentation
        reshaped_effect_variable_values = effect_variable.values.view(-1, 1)

        coefficients, _, _, _ = torch.gels(
            effect_variable_parents_values,
            reshaped_effect_variable_values,
        )

        edge_data["direct_effect"] = coefficients[0][0].item()

        return TestResult(
            u=graph.nodes[nodes[0]],
            v=graph.nodes[nodes[1]],
            action=TestResultAction.UPDATE_EDGE_DIRECTED,
            data=edge_data,
        )
//  [python]
def test(self, nodes: Tuple[str], graph: BaseGraphInterface) -> TestResult:
        """
        Calculate the direct effect of each edge in the graph using multivariate regression.
        :param nodes: list of nodes
        :return: A TestResult with the action to take
        """
        effect_variable = graph.nodes[nodes[1]]
        cause_variable = graph.nodes[nodes[0]]

        if graph.undirected_edge_exists(cause_variable, effect_variable):
            # if the edge is not undirected, we do not need to calculate the direct effect
            return

        edge_data = graph.edge_value(cause_variable, effect_variable)

        all_parents_of_effect_variable = graph.parents_of_node(effect_variable)
        parents_of_effect_variable_without_cause_variable = [
            parent
            for parent in all_parents_of_effect_variable
            if parent != cause_variable
        ]

        effect_variable_parents_values = torch.stack(
            [cause_variable.values]
            + [
                parents_of_effect_variable_without_cause_variable[i].values
                for i in range(len(parents_of_effect_variable_without_cause_variable))
            ],
            dim=1,
        )
        # multivariate regression via tridiagonal reduction and SVD, see torch documentation
        reshaped_effect_variable_values = effect_variable.values.view(-1, 1)

        coefficients = torch.linalg.lstsq(
            effect_variable_parents_values,
            reshaped_effect_variable_values,
            driver="gelsd",
        ).solution

        edge_data["direct_effect"] = coefficients[0][0].item()

        return TestResult(
            u=graph.nodes[nodes[0]],
            v=graph.nodes[nodes[1]],
            action=TestResultAction.UPDATE_EDGE_DIRECTED,
            data=edge_data,
        )
//  [python]
def solve_lstsq(B, c, lmbd=2, verbose=False, method='lstsq'):
    '''Least-squares of form Bx=c
    B: basis, tensor of shape [N,P]
    c: rhs, tensor of shape [N] (target deformation projected onto normal)
    '''
    assert method in ["lstsq", "normalsolve"], "Available methods are lstsq and normalsolve"
    with torch.no_grad():
        ## Driver depending on condition number and device
        ## driver=gels  is fast, but sometimes instable for high condition numbers
        ## driver=gelsd is much more stable, but also much slower
        if verbose: t0 = time()

        if method=='lstsq':
            if lmbd:
                ## Tikhonov regularization
                B_tik = torch.vstack([B, lmbd*torch.eye(len(B.T), device=B.device)])
                c_tik = torch.hstack([c,    torch.zeros(len(B.T), device=B.device)])
                lstsq = torch.solve(B_tik.T @ B_tik, B_tik.T @ c_tik)
            else:
                lstsq = torch.solve(B.T @ B, B.T @ c)
            sol = lstsq.solution
        elif method=='normalsolve':
            '''
            Assemble the normal equation ourselves,
            add regularization onto diagonal
            and torch.linalg.solve the square system.
            Less accurate than lstsq, but less memory requirements.
            For large lmbd they converge to the same.
            '''
            BTB = B.T@B
            BTc = B.T@c
            BTB += lmbd*torch.eye(len(BTB))
            sol = torch.linalg.solve(BTB, BTc)

        if verbose:
            print(f"Solved lstsq in {(time()-t0):.4f} s")
            print(f"Largest sol entry: {sol.abs().max():.2e}")
            approx = B@sol
            mean_res = (approx - c).square().sum()/len(c)
            print(f"Mean residual: {mean_res:.2e}")

        return sol
//  [python]
def solve_lstsq(B, c, lmbd=2, verbose=False, method='lstsq'):
    '''Least-squares of form Bx=c
    B: basis, tensor of shape [N,P]
    c: rhs, tensor of shape [N] (target deformation projected onto normal)
    '''
    assert method in ["lstsq", "normalsolve"], "Available methods are lstsq and normalsolve"
    with torch.no_grad():
        ## Driver depending on condition number and device
        ## driver=gels  is fast, but sometimes instable for high condition numbers
        ## driver=gelsd is much more stable, but also much slower
        if verbose: t0 = time()

        if method=='lstsq':
            if lmbd:
                ## Tikhonov regularization
                B_tik = torch.vstack([B, lmbd*torch.eye(len(B.T), device=B.device)])
                c_tik = torch.hstack([c,    torch.zeros(len(B.T), device=B.device)])
                lstsq = torch.linalg.lstsq(B_tik, c_tik, driver='gels')
            else:
                lstsq = torch.linalg.lstsq(B, c, driver='gels')
            sol = lstsq.solution
        elif method=='normalsolve':
            '''
            Assemble the normal equation ourselves,
            add regularization onto diagonal
            and torch.linalg.solve the square system.
            Less accurate than lstsq, but less memory requirements.
            For large lmbd they converge to the same.
            '''
            BTB = B.T@B
            BTc = B.T@c
            BTB += lmbd*torch.eye(len(BTB))
            sol = torch.linalg.solve(BTB, BTc)

        if verbose:
            print(f"Solved lstsq in {(time()-t0):.4f} s")
            print(f"Largest sol entry: {sol.abs().max():.2e}")
            approx = B@sol
            mean_res = (approx - c).square().sum()/len(c)
            print(f"Mean residual: {mean_res:.2e}")

        return sol
//  [python]
def _test_btrifact(self, cast):
        a = torch.FloatTensor((((1.3722, -0.9020),
                                (1.8849, 1.9169)),
                               ((0.7187, -1.1695),
                                (-0.0139, 1.3572)),
                               ((-1.6181, 0.7148),
                                (1.3728, 0.1319))))
        a = cast(a)
        info = cast(torch.IntTensor())
        a_LU = a.btrifact(info=info)
        self.assertEqual(info.abs().sum(), 0)
        P, a_L, a_U = torch.btriunpack(*a_LU)
        a_ = torch.bmm(P, torch.bmm(a_L, a_U))
        self.assertEqual(a_, a)
//  [python]
def _test_btrifact(self, cast):
        a = torch.FloatTensor((((1.3722, -0.9020),
                                (1.8849, 1.9169)),
                               ((0.7187, -1.1695),
                                (-0.0139, 1.3572)),
                               ((-1.6181, 0.7148),
                                (1.3728, 0.1319))))
        a = cast(a)
        info = cast(torch.IntTensor())
        a_LU = a.btrifact(info=info)
        self.assertEqual(info.abs().sum(), 0)
        P, a_L, a_U = torch.lu_unpack(*a_LU)
        a_ = torch.bmm(P, torch.bmm(a_L, a_U))
        self.assertEqual(a_, a)
//  [python]
def _test_btrifact(self, cast):
        a = torch.FloatTensor((((1.3722, -0.9020),
                                (1.8849, 1.9169)),
                               ((0.7187, -1.1695),
                                (-0.0139, 1.3572)),
                               ((-1.6181, 0.7148),
                                (1.3728, 0.1319))))
        a = cast(a)
        a_LU, pivots = a.btrifact()  # test default info

        # test deprecated info argument
        info = cast(torch.IntTensor())
        with warnings.catch_warnings(record=True):
            a_LU, pivots = a.btrifact(info=info)
        self.assertEqual(info.abs().sum(), 0)

        a_LU_, pivots_, info_ = a.btrifact_with_info()
        self.assertEqual(a_LU, a_LU_)
        self.assertEqual(pivots, pivots_)
        self.assertEqual(info, info_)
        P, a_L, a_U = torch.btriunpack(a_LU, pivots)
        a_ = torch.bmm(P, torch.bmm(a_L, a_U))
        self.assertEqual(a_, a)
//  [python]
def _test_btrifact(self, cast):
        a = torch.FloatTensor((((1.3722, -0.9020),
                                (1.8849, 1.9169)),
                               ((0.7187, -1.1695),
                                (-0.0139, 1.3572)),
                               ((-1.6181, 0.7148),
                                (1.3728, 0.1319))))
        a = cast(a)
        a_LU, pivots = a.btrifact()  # test default info

        # test deprecated info argument
        info = cast(torch.IntTensor())
        with warnings.catch_warnings(record=True):
            a_LU, pivots = a.btrifact(info=info)
        self.assertEqual(info.abs().sum(), 0)

        a_LU_, pivots_, info_ = a.btrifact_with_info()
        self.assertEqual(a_LU, a_LU_)
        self.assertEqual(pivots, pivots_)
        self.assertEqual(info, info_)
        P, a_L, a_U = torch.lu_unpack(a_LU, pivots)
        a_ = torch.bmm(P, torch.bmm(a_L, a_U))
        self.assertEqual(a_, a)
//  [python]
def __init__(self, dim):
        super().__init__()
        self.dim = dim
        Q = torch.nn.init.orthogonal_(torch.randn(dim, dim))
        P, L, U = torch.lu_unpack(*Q.lu())
        self.P = P # remains fixed during optimization
        self.L = nn.Parameter(L) # lower triangular portion
        self.S = nn.Parameter(U.diag()) # "crop out" the diagonal to its own parameter
        self.U = nn.Parameter(torch.triu(U, diagonal=1))
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1ConvLUS, self).__init__()
        # 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]
        p, l, u = torch.lu(W)

        self.register_buffer('p', p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(l, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer('lower_diag', lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(u))
        self.upper = nn.Parameter(torch.triu(u, 1))
        self.cache_inverse = cache_inverse
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1ConvLUS, self).__init__()
        # 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]
        p, lower, upper = torch.lu_unpack(*torch.lu(W))

        self.register_buffer('p', p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(lower, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer('lower_diag', lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(upper))
        self.upper = nn.Parameter(torch.triu(upper, 1))
        self.cache_inverse = cache_inverse
//  [python]
def factor_kkt(S_LU, R, d):
    """ Factor the U22 block that we can only do after we know D. """
    nBatch, nineq = d.size()
    neq = S_LU[1].size(1) - nineq
    # TODO: There's probably a better way to add a batched diagonal.
    global factor_kkt_eye
    if factor_kkt_eye is None or factor_kkt_eye.size() != d.size():
        # print('Updating batchedEye size.')
        factor_kkt_eye = torch.eye(nineq).repeat(
            nBatch, 1, 1).type_as(R).byte()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)

    T_LU = btrifact_hack(T)

    global shown_btrifact_warning
    if shown_btrifact_warning or not T.is_cuda:
        # TODO: Don't use pivoting in most cases because
        # torch.btriunpack is inefficient here:
        oldPivotsPacked = S_LU[1][:, -nineq:] - neq
        oldPivots, _, _ = torch.btriunpack(
            T_LU[0], oldPivotsPacked, unpack_data=False)
        newPivotsPacked = T_LU[1]
        newPivots, _, _ = torch.btriunpack(
            T_LU[0], newPivotsPacked, unpack_data=False)

        # Re-pivot the S_LU_21 block.
        if neq > 0:
            S_LU_21 = S_LU[0][:, -nineq:, :neq]
            S_LU[0][:, -nineq:,
                    :neq] = newPivots.transpose(1, 2).bmm(oldPivots.bmm(S_LU_21))

        # Add the new S_LU_22 block pivots.
        S_LU[1][:, -nineq:] = newPivotsPacked + neq

    # Add the new S_LU_22 block.
    S_LU[0][:, -nineq:, -nineq:] = T_LU[0]
//  [python]
def factor_kkt(S_LU, R, d):
    """ Factor the U22 block that we can only do after we know D. """
    nBatch, nineq = d.size()
    neq = S_LU[1].size(1) - nineq
    # TODO: There's probably a better way to add a batched diagonal.
    global factor_kkt_eye
    if factor_kkt_eye is None or factor_kkt_eye.size() != d.size():
        # print('Updating batchedEye size.')
        factor_kkt_eye = torch.eye(nineq).repeat(
            nBatch, 1, 1).type_as(R).byte()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)

    T_LU = btrifact_hack(T)

    global shown_btrifact_warning
    if shown_btrifact_warning or not T.is_cuda:
        # TODO: Don't use pivoting in most cases because
        # torch.btriunpack is inefficient here:
        oldPivotsPacked = S_LU[1][:, -nineq:] - neq
        oldPivots, _, _ = torch.lu_unpack(
            T_LU[0], oldPivotsPacked, unpack_data=False)
        newPivotsPacked = T_LU[1]
        newPivots, _, _ = torch.lu_unpack(
            T_LU[0], newPivotsPacked, unpack_data=False)

        # Re-pivot the S_LU_21 block.
        if neq > 0:
            S_LU_21 = S_LU[0][:, -nineq:, :neq]
            S_LU[0][:, -nineq:,
                    :neq] = newPivots.transpose(1, 2).bmm(oldPivots.bmm(S_LU_21))

        # Add the new S_LU_22 block pivots.
        S_LU[1][:, -nineq:] = newPivotsPacked + neq

    # Add the new S_LU_22 block.
    S_LU[0][:, -nineq:, -nineq:] = T_LU[0]
//  [python]
def sample_inputs_linalg_det_singular(op_info, device, dtype, requires_grad, **kwargs):
    make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)

    def make_singular_matrix_batch_base(size, rank):
        assert size[-1] == size[-2]
        assert rank > 0 and rank < size[-1]

        n = size[-1]
        a = make_arg(size[:-2] + (n, rank)) / 10
        b = make_arg(size[:-2] + (rank, n)) / 10
        x = a @ b
        lu, pivs, _ = torch.linalg.lu_factor_ex(x)
        p, l, u = torch.lu_unpack(lu, pivs)
        u_diag_abs = u.diagonal(0, -2, -1).abs()
        u_diag_abs_largest = u_diag_abs.max(dim=-1, keepdim=True).values
        u_diag_abs_smallest_idxs = torch.topk(u_diag_abs, k=(n - rank), largest=False).indices
        u.diagonal(0, -2, -1).div_(u_diag_abs_largest)
        u[..., u_diag_abs_smallest_idxs] = torch.finfo(dtype).eps
        matrix = p @ l @ u

        matrix.requires_grad_(requires_grad)
        return matrix

    def sample_generator():
        for batch, size in product(((), (2,), (2, 2)), range(6)):
            shape = batch + (size, size)
            for rank in range(1, size):
                yield make_singular_matrix_batch_base(shape, rank)

    return [SampleInput(t) for t in sample_generator()]
//  [python]
def factor_kkt(S_LU, R, d):
    """ Factor the U22 block that we can only do after we know D. """
    nBatch, nineq = d.size()
    neq = S_LU[1].size(1) - nineq
    # TODO There's probably a better way to add a batched diagonal.
    global factor_kkt_eye
    if factor_kkt_eye is None or factor_kkt_eye.size() != d.size():
        # print('Updating batchedEye size.')
        factor_kkt_eye = torch.eye(nineq).repeat(
            nBatch, 1, 1).type_as(R).byte()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).view(-1)

    T_LU = btrifact_hack(T)

    global shown_btrifact_warning
    if shown_btrifact_warning or not T.is_cuda:
        # TODO Don't use pivoting in most cases because
        # torch.btriunpack is inefficient here:
        oldPivotsPacked = S_LU[1][:, -nineq:] - neq
        oldPivots, _, _ = torch.btriunpack(
            T_LU[0], oldPivotsPacked, unpack_data=False)
        newPivotsPacked = T_LU[1]
        newPivots, _, _ = torch.btriunpack(
            T_LU[0], newPivotsPacked, unpack_data=False)

        # Re-pivot the S_LU_21 block.
        if neq > 0:
            S_LU_21 = S_LU[0][:, -nineq:, :neq]
            S_LU[0][:, -nineq:,
                    :neq] = newPivots.transpose(1, 2).bmm(oldPivots.bmm(S_LU_21))

        # Add the new S_LU_22 block pivots.
        S_LU[1][:, -nineq:] = newPivotsPacked + neq

    # Add the new S_LU_22 block.
    S_LU[0][:, -nineq:, -nineq:] = T_LU[0]
//  [python]
def factor_kkt(S_LU, R, d):
    """ Factor the U22 block that we can only do after we know D. """
    nBatch, nineq = d.size()
    neq = S_LU[1].size(1) - nineq
    # TODO There's probably a better way to add a batched diagonal.
    global factor_kkt_eye
    if factor_kkt_eye is None or factor_kkt_eye.size() != d.size():
        # print('Updating batchedEye size.')
        factor_kkt_eye = torch.eye(nineq).repeat(
            nBatch, 1, 1).type_as(R).byte()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).view(-1)

    T_LU = btrifact_hack(T)

    global shown_btrifact_warning
    if shown_btrifact_warning or not T.is_cuda:
        # TODO Don't use pivoting in most cases because
        # torch.btriunpack is inefficient here:
        oldPivotsPacked = S_LU[1][:, -nineq:] - neq
        oldPivots, _ = torch.lu_unpack(
            T_LU[0], oldPivotsPacked, unpack_data=False)
        newPivotsPacked = T_LU[1]
        newPivots, _ = torch.lu_unpack(
            T_LU[0], newPivotsPacked, unpack_data=False)

        # Re-pivot the S_LU_21 block.
        if neq > 0:
            S_LU_21 = S_LU[0][:, -nineq:, :neq]
            S_LU[0][:, -nineq:,
                    :neq] = newPivots.transpose(1, 2).bmm(oldPivots.bmm(S_LU_21))

        # Add the new S_LU_22 block pivots.
        S_LU[1][:, -nineq:] = newPivotsPacked + neq

    # Add the new S_LU_22 block.
    S_LU[0][:, -nineq:, -nineq:] = T_LU[0]
//  [python]
def pre_factor_kkt(Q, G, A):
    """ Perform all one-time factorizations and cache relevant matrix products"""
    nineq, nz, neq, nBatch = get_sizes(G, A)

    try:
        Q_LU = btrifact_hack(Q)
    except:
        raise RuntimeError("""
qpth Error: Cannot perform LU factorization on Q.
Please make sure that your Q matrix is PSD and has
a non-zero diagonal.
""")

    # S = [ A Q^{-1} A^T        A Q^{-1} G^T          ]
    #     [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    #
    # We compute a partial LU decomposition of the S matrix
    # that can be completed once D^{-1} is known.
    # See the 'Block LU factorization' part of our website
    # for more details.

    G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU))
    R = G_invQ_GT.clone()
    S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \
        .repeat(nBatch, 1).type_as(Q).int()
    if neq > 0:
        invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU)
        A_invQ_AT = torch.bmm(A, invQ_AT)
        G_invQ_AT = torch.bmm(G, invQ_AT)

        LU_A_invQ_AT = btrifact_hack(A_invQ_AT)
        P_A_invQ_AT, L_A_invQ_AT, U_A_invQ_AT = torch.btriunpack(*LU_A_invQ_AT)
        P_A_invQ_AT = P_A_invQ_AT.type_as(A_invQ_AT)

        S_LU_11 = LU_A_invQ_AT[0]
        U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT)
                           ).btrisolve(*LU_A_invQ_AT)
        S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv)
        T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT)
        S_LU_12 = U_A_invQ_AT.bmm(T)
        S_LU_22 = torch.zeros(nBatch, nineq, nineq).type_as(Q)
        S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2),
                               torch.cat((S_LU_21, S_LU_22), 2)),
                              1)
        S_LU_pivots[:, :neq] = LU_A_invQ_AT[1]

        R -= G_invQ_AT.bmm(T)
    else:
        S_LU_data = torch.zeros(nBatch, nineq, nineq).type_as(Q)

    S_LU = [S_LU_data, S_LU_pivots]
    return Q_LU, S_LU, R
//  [python]
def pre_factor_kkt(Q, G, A):
    """ Perform all one-time factorizations and cache relevant matrix products"""
    nineq, nz, neq, nBatch = get_sizes(G, A)

    try:
        Q_LU = btrifact_hack(Q)
    except:
        raise RuntimeError("""
qpth Error: Cannot perform LU factorization on Q.
Please make sure that your Q matrix is PSD and has
a non-zero diagonal.
""")

    # S = [ A Q^{-1} A^T        A Q^{-1} G^T          ]
    #     [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    #
    # We compute a partial LU decomposition of the S matrix
    # that can be completed once D^{-1} is known.
    # See the 'Block LU factorization' part of our website
    # for more details.

    G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU))
    R = G_invQ_GT.clone()
    S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \
        .repeat(nBatch, 1).type_as(Q).int()
    if neq > 0:
        invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU)
        A_invQ_AT = torch.bmm(A, invQ_AT)
        G_invQ_AT = torch.bmm(G, invQ_AT)

        LU_A_invQ_AT = btrifact_hack(A_invQ_AT)
        P_A_invQ_AT, L_A_invQ_AT, U_A_invQ_AT = torch.lu_unpack(*LU_A_invQ_AT)
        P_A_invQ_AT = P_A_invQ_AT.type_as(A_invQ_AT)

        S_LU_11 = LU_A_invQ_AT[0]
        U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT)
                           ).btrisolve(*LU_A_invQ_AT)
        S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv)
        T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT)
        S_LU_12 = U_A_invQ_AT.bmm(T)
        S_LU_22 = torch.zeros(nBatch, nineq, nineq).type_as(Q)
        S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2),
                               torch.cat((S_LU_21, S_LU_22), 2)),
                              1)
        S_LU_pivots[:, :neq] = LU_A_invQ_AT[1]

        R -= G_invQ_AT.bmm(T)
    else:
        S_LU_data = torch.zeros(nBatch, nineq, nineq).type_as(Q)

    S_LU = [S_LU_data, S_LU_pivots]
    return Q_LU, S_LU, R
//  [python]
def sample_inputs_linalg_det_singular(op_info, device, dtype, requires_grad, **kwargs):
    make_arg = partial(make_tensor, device=device, dtype=dtype)

    def make_singular_matrix_batch_base(size, rank):
        assert size[-1] == size[-2]
        assert rank > 0 and rank < size[-1]

        n = size[-1]
        a = make_arg(size[:-2] + (n, rank)) / 10
        b = make_arg(size[:-2] + (rank, n)) / 10
        x = a @ b
        lu, pivs, _ = torch.linalg.lu_factor_ex(x)
        p, l, u = torch.lu_unpack(lu, pivs)
        u_diag_abs = u.diagonal(0, -2, -1).abs()
        u_diag_abs_largest = u_diag_abs.max(dim=-1, keepdim=True).values
        u_diag_abs_smallest_idxs = torch.topk(u_diag_abs, k=(n - rank), largest=False).indices
        u.diagonal(0, -2, -1).div_(u_diag_abs_largest)
        u.diagonal(0, -2, -1)[..., u_diag_abs_smallest_idxs] = torch.finfo(dtype).eps
        matrix = p @ l @ u

        matrix.requires_grad_(requires_grad)
        return matrix

    def sample_generator():
        for batch, size in product(((), (2,), (2, 2)), range(6)):
            shape = batch + (size, size)
            for rank in range(1, size):
                yield make_singular_matrix_batch_base(shape, rank)

    return [SampleInput(t) for t in sample_generator()]
//  [python]
def __init__(self, dim):
        super().__init__()
        self.dim = dim
        Q = nn.init.orthogonal_(torch.randn(dim, dim).to(device))
        P, L, U = torch.lu_unpack(*torch.lu(Q))
        self.register_buffer('P', P)
        self.L = nn.Parameter(L) # lower triangular portion
        self.S = nn.Parameter(U.diag()) # "crop out" the diagonal to its own parameter
        self.U = nn.Parameter(torch.triu(U, diagonal=1))
//  [python]
def __init__(self, dim):
        super(Invertible1x1Conv, self).__init__()
        self.dim = dim
        Q = torch.nn.init.orthogonal_(torch.randn(dim, dim))
        P, L, U = torch.lu_unpack(*Q.lu())
        self.P = P  # remains fixed during optimization
        self.L = nn.Parameter(L)  # lower triangular portion
        self.S = nn.Parameter(U.diag())  # "crop out" the diagonal to its own parameter
        self.U = nn.Parameter(torch.triu(U, diagonal=1))
//  [python]
def pre_factor_kkt(Q, G, F, A):
    """ Perform all one-time factorizations and cache relevant matrix products"""
    nineq, nz, neq, nBatch = get_sizes(G, A)

    try:
        Q_LU = btrifact_hack(Q)
    except:
        raise RuntimeError("""
lcp Error: Cannot perform LU factorization on Q.
Please make sure that your Q matrix is PSD and has
a non-zero diagonal.
""")

    # S = [ A Q^{-1} A^T        A Q^{-1} G^T          ]
    #     [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    #
    # We compute a partial LU decomposition of the S matrix
    # that can be completed once D^{-1} is known.
    # See the 'Block LU factorization' part of our website
    # for more details.

    G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU)) + F
    R = G_invQ_GT.clone()
    S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \
        .repeat(nBatch, 1).type_as(Q).int()
    if neq > 0:
        invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU)
        A_invQ_AT = torch.bmm(A, invQ_AT)
        G_invQ_AT = torch.bmm(G, invQ_AT)

        LU_A_invQ_AT = btrifact_hack(A_invQ_AT)
        P_A_invQ_AT, L_A_invQ_AT, U_A_invQ_AT = torch.btriunpack(*LU_A_invQ_AT)
        P_A_invQ_AT = P_A_invQ_AT.type_as(A_invQ_AT)

        S_LU_11 = LU_A_invQ_AT[0]
        U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT)
                           ).btrisolve(*LU_A_invQ_AT)
        S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv)
        T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT)
        S_LU_12 = U_A_invQ_AT.bmm(T)
        S_LU_22 = torch.zeros(nBatch, nineq, nineq).type_as(Q)
        S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2),
                               torch.cat((S_LU_21, S_LU_22), 2)),
                              1)
        S_LU_pivots[:, :neq] = LU_A_invQ_AT[1]

        R -= G_invQ_AT.bmm(T)
    else:
        S_LU_data = torch.zeros(nBatch, nineq, nineq).type_as(Q)

    S_LU = [S_LU_data, S_LU_pivots]
    return Q_LU, S_LU, R
//  [python]
def pre_factor_kkt(Q, G, F, A):
    """ Perform all one-time factorizations and cache relevant matrix products"""
    nineq, nz, neq, nBatch = get_sizes(G, A)

    try:
        Q_LU = btrifact_hack(Q)
    except:
        raise RuntimeError("""
lcp Error: Cannot perform LU factorization on Q.
Please make sure that your Q matrix is PSD and has
a non-zero diagonal.
""")

    # S = [ A Q^{-1} A^T        A Q^{-1} G^T          ]
    #     [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    #
    # We compute a partial LU decomposition of the S matrix
    # that can be completed once D^{-1} is known.
    # See the 'Block LU factorization' part of our website
    # for more details.

    G_invQ_GT = torch.bmm(G, G.transpose(1, 2).btrisolve(*Q_LU)) + F
    R = G_invQ_GT.clone()
    S_LU_pivots = torch.IntTensor(range(1, 1 + neq + nineq)).unsqueeze(0) \
        .repeat(nBatch, 1).type_as(Q).int()
    if neq > 0:
        invQ_AT = A.transpose(1, 2).btrisolve(*Q_LU)
        A_invQ_AT = torch.bmm(A, invQ_AT)
        G_invQ_AT = torch.bmm(G, invQ_AT)

        LU_A_invQ_AT = btrifact_hack(A_invQ_AT)
        P_A_invQ_AT, L_A_invQ_AT, U_A_invQ_AT = torch.lu_unpack(*LU_A_invQ_AT)
        P_A_invQ_AT = P_A_invQ_AT.type_as(A_invQ_AT)

        S_LU_11 = LU_A_invQ_AT[0]
        U_A_invQ_AT_inv = (P_A_invQ_AT.bmm(L_A_invQ_AT)
                           ).btrisolve(*LU_A_invQ_AT)
        S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv)
        T = G_invQ_AT.transpose(1, 2).btrisolve(*LU_A_invQ_AT)
        S_LU_12 = U_A_invQ_AT.bmm(T)
        S_LU_22 = torch.zeros(nBatch, nineq, nineq).type_as(Q)
        S_LU_data = torch.cat((torch.cat((S_LU_11, S_LU_12), 2),
                               torch.cat((S_LU_21, S_LU_22), 2)),
                              1)
        S_LU_pivots[:, :neq] = LU_A_invQ_AT[1]

        R -= G_invQ_AT.bmm(T)
    else:
        S_LU_data = torch.zeros(nBatch, nineq, nineq).type_as(Q)

    S_LU = [S_LU_data, S_LU_pivots]
    return Q_LU, S_LU, R
//  [python]
def __init__(self, dim, device='cpu', condition_size=0):
        super().__init__()
        self.conditional = True  # forward backward unchanged when conditional, thus always True
        self.cond_size = condition_size  # for compatibility with the rest of the flow who have this attribute
        self.dim = dim
        self.device = device
        Q = torch.nn.init.orthogonal_(torch.randn(dim, dim).to(self.device))
        P, L, U = torch.lu_unpack(*Q.lu())
        self.P = P.to(self.device)  # remains fixed during optimization
        self.L = nn.Parameter(L).to(self.device)  # lower triangular portion
        self.S = nn.Parameter(U.diag()).to(self.device)  # "crop out" the diagonal to its own parameter
        self.U = nn.Parameter(torch.triu(U, diagonal=1)).to(self.device)
//  [python]
def __init__(self, x_shape, num_u_channels=0):
        super().__init__(x_shape, num_u_channels)

        P, L, U = torch.lu(self.weights_init)
        s = torch.diag(U)
        log_s = torch.log(torch.abs(s))
        upper = torch.triu(U, 1)

        self.register_buffer('P', P)
        self.register_buffer('sign_s', torch.sign(s))
        self.register_buffer('l_mask', torch.tril(torch.ones(self.weight_shape), -1))
        self.register_buffer('eye', torch.eye(*self.weight_shape))

        self.lower = nn.Parameter(L)
        self.log_s = nn.Parameter(log_s)
        self.upper = nn.Parameter(upper)

        self.bias = nn.Parameter(torch.zeros(x_shape[0], *x_shape[1:]))
//  [python]
def __init__(self, x_shape, num_u_channels=0):
        super().__init__(x_shape, num_u_channels)

        P, lower, upper = torch.lu_unpack(*torch.lu(self.weights_init))
        s = torch.diag(upper)
        log_s = torch.log(torch.abs(s))
        upper = torch.triu(upper, 1)

        self.register_buffer('P', P)
        self.register_buffer('sign_s', torch.sign(s))
        self.register_buffer('l_mask', torch.tril(torch.ones(self.weight_shape), -1))
        self.register_buffer('eye', torch.eye(*self.weight_shape))

        self.lower = nn.Parameter(lower)
        self.log_s = nn.Parameter(log_s)
        self.upper = nn.Parameter(upper)

        self.bias = nn.Parameter(torch.zeros(x_shape[0], *x_shape[1:]))
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1ConvLUS, self).__init__()
        # 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]
        p, l, u = torch.lu(W)

        self.register_buffer("p", p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(l, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer("lower_diag", lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(u))
        self.upper = nn.Parameter(torch.triu(u, 1))
        self.cache_inverse = cache_inverse
//  [python]
def __init__(self, c, cache_inverse=False):
        super(Invertible1x1ConvLUS, self).__init__()
        # 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]
        p, lower, upper = torch.lu_unpack(*torch.lu(W))

        self.register_buffer("p", p)
        # diagonals of lower will always be 1s anyway
        lower = torch.tril(lower, -1)
        lower_diag = torch.diag(torch.eye(c, c))
        self.register_buffer("lower_diag", lower_diag)
        self.lower = nn.Parameter(lower)
        self.upper_diag = nn.Parameter(torch.diag(upper))
        self.upper = nn.Parameter(torch.triu(upper, 1))
        self.cache_inverse = cache_inverse