jiahuimbzuai/CoQuIR huggingface dataset:
//  [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).bool()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)

    T_LU = lu_hack(T)

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

        # 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).bool()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)

    T_LU = lu_hack(T)

    if not T.is_cuda:
        # TODO: Don't use pivoting in most cases because
        # torch.lu_unpack 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 multiplier(self,rho):
        # m = (A^T*A+Φ^Τ*Φ)^-1
        # Instead of calculating directly the inverse, we take the LU factorization of A^T*A+Φ^Τ*Φ

        ata = torch.mm(self.a.t(), self.a)
        ftf = torch.mm(self.phi.t(),self.phi)
        m = ata + rho * ftf
        m_lu, pivots = m.lu()
        L, U = torch.lu_unpack(m_lu, pivots)
        Linv = torch.linalg.inv(L)
        Uinv = torch.linalg.inv(U)

        return Linv, Uinv
//  [python]
def multiplier(self,rho):
        # m = (A^T*A+Φ^Τ*Φ)^-1
        # Instead of calculating directly the inverse, we take the LU factorization of A^T*A+Φ^Τ*Φ

        ata = torch.mm(self.a.t(), self.a)
        ftf = torch.mm(self.phi.t(),self.phi)
        m = ata + rho * ftf
        m_lu, _ = m.lu()
        _, L, U = torch.lu_unpack(m_lu, _)
        Linv = torch.linalg.inv(L)
        Uinv = torch.linalg.inv(U)

        return Linv, Uinv
//  [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 = lu_hack(T)

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

        # 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 = lu_hack(T)

    if not T.is_cuda:
        # TODO: Don't use pivoting in most cases because
        # torch.lu_unpack 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 = lu_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).lu_solve(*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).lu_solve(*Q_LU)
        A_invQ_AT = torch.bmm(A, invQ_AT)
        G_invQ_AT = torch.bmm(G, invQ_AT)

        LU_A_invQ_AT = lu_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)
                           ).lu_solve(*LU_A_invQ_AT)
        S_LU_21 = G_invQ_AT.bmm(U_A_invQ_AT_inv)
        T = G_invQ_AT.transpose(1, 2).lu_solve(*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 _gen_shape_inputs_linalg_triangular_solve(self, shape, dtype, device, well_conditioned=False):
        make_arg = partial(make_tensor, dtype=dtype, device=device)
        make_randn = partial(torch.randn, dtype=dtype, device=device)
        b, n, k = shape
        for left, uni, expand_a, tr_a, conj_a, expand_b, tr_b, conj_b in product((True, False), repeat=8):
            # expand means that we generate a batch of matrices with a stride of zero in the batch dimension
            if (conj_a or conj_b) and not dtype.is_complex:
                continue
            # We just expand on the batch size
            if (expand_a or expand_b) and b == 1:
                continue

            size_a = (b, n, n) if left else (b, k, k)
            size_b = (b, n, k) if not tr_b else (b, k, n)

            # If expand_a or expand_b, we'll expand them to the correct size later
            if b == 1 or expand_a:
                size_a = size_a[1:]
            if b == 1 or expand_b:
                size_b = size_b[1:]

            if well_conditioned:
                P, L, U = torch.lu(make_randn(*size_a))
                if uni:
                    # A = L from PLU
                    A = L.transpose(-2, -1).contiguous()
                else:
                    # A = U from PLU
                    A = U.contiguous()
            else:
                A = make_arg(size_a)
                A.triu_()

            diag = A.diagonal(0, -2, -1)
            if uni:
                diag.fill_(1.)
            else:
                diag[diag.abs() < 1e-6] = 1.

            B = make_arg(size_b)

            if tr_a:
                A.transpose_(-2, -1)
            if tr_b:
                B.transpose_(-2, -1)
            if conj_a:
                A = A.conj()
            if conj_b:
                B = B.conj()
            if expand_a:
                A = A.expand(b, *size_a)
            if expand_b:
                B = B.expand(b, n, k)
            yield A, B, left, not tr_a, uni
//  [python]
def _gen_shape_inputs_linalg_triangular_solve(self, shape, dtype, device, well_conditioned=False):
        make_arg = partial(make_tensor, dtype=dtype, device=device)
        make_randn = partial(torch.randn, dtype=dtype, device=device)
        b, n, k = shape
        for left, uni, expand_a, tr_a, conj_a, expand_b, tr_b, conj_b in product((True, False), repeat=8):
            # expand means that we generate a batch of matrices with a stride of zero in the batch dimension
            if (conj_a or conj_b) and not dtype.is_complex:
                continue
            # We just expand on the batch size
            if (expand_a or expand_b) and b == 1:
                continue

            size_a = (b, n, n) if left else (b, k, k)
            size_b = (b, n, k) if not tr_b else (b, k, n)

            # If expand_a or expand_b, we'll expand them to the correct size later
            if b == 1 or expand_a:
                size_a = size_a[1:]
            if b == 1 or expand_b:
                size_b = size_b[1:]

            if well_conditioned:
                PLU = torch.lu_unpack(*torch.lu(make_randn(*size_a)))
                if uni:
                    # A = L from PLU
                    A = PLU[1].transpose(-2, -1).contiguous()
                else:
                    # A = U from PLU
                    A = PLU[2].contiguous()
            else:
                A = make_arg(size_a)
                A.triu_()

            diag = A.diagonal(0, -2, -1)
            if uni:
                diag.fill_(1.)
            else:
                diag[diag.abs() < 1e-6] = 1.

            B = make_arg(size_b)

            if tr_a:
                A.transpose_(-2, -1)
            if tr_b:
                B.transpose_(-2, -1)
            if conj_a:
                A = A.conj()
            if conj_b:
                B = B.conj()
            if expand_a:
                A = A.expand(b, *size_a)
            if expand_b:
                B = B.expand(b, n, k)
            yield A, B, left, not tr_a, uni
//  [python]
def multiplier(self,rho):
        # m = (A^T*A+Φ^Τ*Φ)^-1
        # Instead of calculating directly the inverse, we take the LU factorization of A^T*A+Φ^Τ*Φ

        ata = torch.mm(self.a.t(), self.a) 
        ftf = torch.mm(self.phi.t(), self.phi)
        m = ata + self.rho * ftf
        m_lu, pivots = m.lu()
        L, U = torch.lu_unpack(m_lu, pivots)
        Linv = torch.linalg.inv(L)
        Uinv = torch.linalg.inv(U)

        return Linv, Uinv
//  [python]
def multiplier(self,rho):
        # m = (A^T*A+Φ^Τ*Φ)^-1
        # Instead of calculating directly the inverse, we take the LU factorization of A^T*A+Φ^Τ*Φ

        ata = torch.mm(self.a.t(), self.a) 
        ftf = torch.mm(self.phi.t(), self.phi)
        m = ata + self.rho * ftf
        m_lu, _ = m.lu()
        _, L, U = torch.lu_unpack(m_lu, _)
        Linv = torch.linalg.inv(L)
        Uinv = torch.linalg.inv(U)

        return Linv, Uinv
//  [python]
def __init__(self, c, cache_inverse=False):
        super(DataInitializedInvertible1x1Conv, self).__init__()
        self.register_buffer('input_mean', torch.zeros(c, 1))
        self.register_buffer('initialized', torch.tensor(False))
        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
        # same initialization as before, though we'll override these values soon anyway
        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
        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(DataInitializedInvertible1x1Conv, self).__init__()
        self.register_buffer('input_mean', torch.zeros(c, 1))
        self.register_buffer('initialized', torch.tensor(False))
        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]
        # same initialization as before, though we'll override these values soon anyway
        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
        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.lu_unpack(
            T_LU[0], oldPivotsPacked, unpack_data=False)[0]
        newPivotsPacked = T_LU[1]
        newPivots = torch.lu_unpack(
            T_LU[0], newPivotsPacked, unpack_data=False)[0]

        # 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 __init__(self, dim):
        super().__init__()
        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
        # "crop out" diagonal, stored in S
        self.U = nn.Parameter(torch.triu(U, diagonal=1))
//  [python]
def test_linalg_lu_factor_and_lu(self, device, dtype):
        # Tests lu, linalg.lu_factor and linalg.lu_factor_ex
        from torch.testing._internal.common_utils import random_matrix

        def run_test(A, pivot, singular, fn):
            k = min(A.shape[-2:])
            batch = A.shape[:-2]
            check_errors = (fn == torch.linalg.lu_factor)
            if singular and check_errors:
                # It may or may not throw as the LU decomposition without pivoting
                # may still succeed for singular matrices
                try:
                    LU, pivots = fn(A, pivot=pivot)
                except RuntimeError:
                    return
            else:
                LU, pivots = fn(A, pivot=pivot)[:2]

            self.assertEqual(LU.size(), A.shape)
            self.assertEqual(pivots.size(), batch + (k,))

            if not pivot:
                self.assertEqual(pivots, torch.arange(1, 1 + k, device=device, dtype=torch.int32).expand(batch + (k, )))

            P, L, U = torch.lu_unpack(LU, pivots)

            self.assertEqual(P @ L @ U, A)

        sizes = ((3, 3), (5, 5), (4, 2), (3, 4), (0, 0), (0, 1), (1, 0))
        batches = ((0,), (2,), (3,), (1, 0), (3, 5))
        # Non pivoting just implemented for CUDA
        pivots = (True, False) if self.device_type == "cuda" else (True,)
        fns = (partial(torch.lu, get_infos=True), torch.linalg.lu_factor, torch.linalg.lu_factor_ex)
        for ms, batch, pivot, singular, fn in itertools.product(sizes, batches, pivots, (True, False), fns):
            m, n = ms
            A = random_matrix(m, n, *batch, singular=singular, dtype=dtype, device=device)
            # Just do one of them on singular matrices
            if A.numel() == 0 and not singular:
                continue
            run_test(A, pivot, singular, fn)

            # Reproducer of a magma bug,
            # see https://bitbucket.org/icl/magma/issues/13/getrf_batched-kernel-produces-nans-on
            # This is also a bug in cuSOLVER < 11.3
            if (dtype == torch.double
               and singular
               and (torch.version.cuda is None or
                    torch.version.cuda.split('.') >= ["11", "3"])):
                A = torch.ones(batch + ms, dtype=dtype, device=device)
                run_test(A, pivot, singular, fn)

        # Info should be positive for rank deficient matrices
        A = torch.ones(5, 3, 3, device=device)
        self.assertTrue((torch.linalg.lu_factor_ex(A, pivot=True).info >= 0).all())

        if self.device_type == 'cpu':
            # Error checking, no pivoting variant on CPU
            with self.assertRaisesRegex(RuntimeError, 'LU without pivoting is not implemented on the CPU'):
                torch.lu(torch.empty(1, 2, 2), pivot=False)

            with self.assertRaisesRegex(RuntimeError, 'LU without pivoting is not implemented on the CPU'):
                torch.linalg.lu_factor(torch.empty(1, 2, 2), pivot=False)
//  [python]
def sample_inputs_linalg_det_singular(op_info, device, dtype, requires_grad):
    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]

        with torch.no_grad():
            n = size[-1]
            a = make_arg(size[:-2] + (n, rank)) / 10
            b = make_arg(size[:-2] + (rank, n)) / 10

            x = a @ b
            lu, pivs = x.lu()
            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

        assert (matrix.det().abs() < torch.finfo(dtype).eps * torch.linalg.matrix_norm(matrix)).all().item()

        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 getl(a_lu, pivots):
    a_lu_float = a_lu.float()
    a_lu_torch = torch.tensor(a_lu_float, requires_grad=True)
    out_torch = torch.lu_unpack(a_lu_torch, pivots)
    out_torch[1].backward(gradient=out_torch[1])
    return a_lu_torch.grad
//  [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 getu(a_lu, pivots):
    a_lu_float = a_lu.float()
    a_lu_torch = torch.tensor(a_lu_float, requires_grad=True)
    out_torch = torch.lu_unpack(a_lu_torch, pivots)
    out_torch[2].backward(gradient=out_torch[2])
    return a_lu_torch.grad
//  [python]
def forward(ctx, input):

        # LUP decompose the matrices
        inp_lu, pivots = input.lu()
        perm, inpl, inpu = torch.lu_unpack(inp_lu, pivots)

        # get the number of permuations
        s = (pivots != torch.as_tensor(
            range(1, input.shape[1]+1)).int()).sum(1).type(torch.get_default_dtype())

        # get the prod of the diag of U
        d = torch.diagonal(inpu, dim1=-2, dim2=-1).prod(1)

        # assemble
        det = ((-1)**s * d)
        ctx.save_for_backward(input, det)

        return det
//  [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).bool()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)


    T_LU = lu_hack(T)


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

        # 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).bool()
    T = R.clone()
    T[factor_kkt_eye] += (1. / d).squeeze().view(-1)


    T_LU = lu_hack(T)


    if not T.is_cuda:
        # TODO: Don't use pivoting in most cases because
        # torch.lu_unpack 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 test_lu_unpack_grad_graph_float():
    """
    Feature: LuUnpackGrad gpu TEST.
    Description: 4d - float32 test case for LuUnpackGrad
    Expectation: the result match to numpy
    """
    a = torch.randn(3, 3, 3, 3)
    loss = 1e-5
    a_lu, pivots = a.lu()
    out = torch.lu_unpack(a_lu, pivots)

    a_l_torch = getl(a_lu, pivots)
    a_u_torch = getu(a_lu, pivots)

    a_l_mindspore = Tensor(out[1].numpy(), mstype.float32)
    a_u_mindspore = Tensor(out[2].numpy(), mstype.float32)
    a_lu_mindspore = Tensor(a_lu.numpy(), mstype.float32)

    net = NetLuUnpackGrad(l_grad_flag=True, u_grad_flag=True)
    l_grad_mindspore, u_grad_mindspore = net(a_l_mindspore, a_u_mindspore, a_lu_mindspore)

    assert np.allclose(a_l_torch, torch.tensor(l_grad_mindspore.asnumpy()), atol=loss)
    assert np.allclose(a_u_torch, torch.tensor(u_grad_mindspore.asnumpy()), atol=loss)
//  [python]
def test_lu_unpack_grad_pynative_float_error():
    """
    Feature: LuUnpackGrad gpu TEST.
    Description: 4d - uint16 test the unsupported type for LuUnpackGrad
    Expectation: the result match to numpy
    """
    a = torch.randn(8, 7, 4, 2)
    a_lu, pivots = a.lu()
    out = torch.lu_unpack(a_lu, pivots)

    a_l_mindspore = Tensor(out[1].numpy(), mstype.uint16)
    a_u_mindspore = Tensor(out[2].numpy(), mstype.uint16)
    a_lu_mindspore = Tensor(a_lu.numpy(), mstype.uint16)

    with pytest.raises(TypeError):
        net = NetLuUnpackGrad(l_grad_flag=True, u_grad_flag=True)
        net(a_l_mindspore, a_u_mindspore, a_lu_mindspore)
//  [python]
def __init__(self, dim_inputs):
        super().__init__()
        self.dim_inputs = dim_inputs
        w_shape = [dim_inputs, dim_inputs]
        w_init = torch.qr(torch.randn(w_shape))[0]

        p, l, u = torch.lu(w_init)
        s = torch.diag(u)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        u = torch.triu(u, 1)
        l_mask = torch.tril(torch.ones(w_shape), -1)
        eye = torch.eye(*w_shape)

        self.register_buffer('p', p)
        self.register_buffer('sign_s', sign_s)
        self.lower = nn.Parameter(l)
        self.log_s = nn.Parameter(log_s)
        self.upper = nn.Parameter(u)
        self.register_buffer('l_mask', l_mask)
        self.register_buffer('eye', eye)
//  [python]
def __init__(self, dim_inputs):
        super().__init__()
        self.dim_inputs = dim_inputs
        w_shape = [dim_inputs, dim_inputs]
        w_init = torch.qr(torch.randn(w_shape))[0]

        p, lower, upper = torch.lu_unpack(*torch.lu(w_init))
        s = torch.diag(upper)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        upper = torch.triu(upper, 1)
        l_mask = torch.tril(torch.ones(w_shape), -1)
        eye = torch.eye(*w_shape)

        self.register_buffer('p', p)
        self.register_buffer('sign_s', sign_s)
        self.lower = nn.Parameter(lower)
        self.log_s = nn.Parameter(log_s)
        self.upper = nn.Parameter(upper)
        self.register_buffer('l_mask', l_mask)
        self.register_buffer('eye', eye)
//  [python]
def test_lu_unpack_grad_pynative_float():
    """
    Feature: LuUnpackGrad gpu TEST.
    Description: 4d - float32 test case for LuUnpackGrad
    Expectation: the result match to numpy
    """
    a = torch.randn(5, 4, 3, 2)
    loss = 1e-5
    a_lu, pivots = a.lu()
    out = torch.lu_unpack(a_lu, pivots)

    a_l_torch = getl(a_lu, pivots)
    a_u_torch = getu(a_lu, pivots)

    a_l_mindspore = Tensor(out[1].numpy(), mstype.float32)
    a_u_mindspore = Tensor(out[2].numpy(), mstype.float32)
    a_lu_mindspore = Tensor(a_lu.numpy(), mstype.float32)

    net = NetLuUnpackGrad(l_grad_flag=True, u_grad_flag=True)
    l_grad_mindspore, u_grad_mindspore = net(a_l_mindspore, a_u_mindspore, a_lu_mindspore)

    assert np.allclose(a_l_torch, torch.tensor(l_grad_mindspore.asnumpy()), atol=loss)
    assert np.allclose(a_u_torch, torch.tensor(u_grad_mindspore.asnumpy()), atol=loss)
//  [python]
def __init__(self, inCh):
        super(InvConvLU1x1, self).__init__()
        w_shape = [inCh, inCh]
        w_init = torch.qr(torch.randn(*w_shape))[0]

        p, l, u = torch.lu(w_init)
        s = torch.diag(u)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        u = torch.triu(u, 1)
        l_mask = torch.tril(torch.ones(w_shape), -1)
        eye = torch.eye(*w_shape)

        self.register_buffer('p', p)  # .cuda() will work only on register_buffer
        self.register_buffer('sign_s', sign_s)
        self.register_buffer('l_mask', l_mask)
        self.register_buffer('eye', eye)

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

        self.w_shape = w_shape
//  [python]
def __init__(self, inCh):
        super(InvConvLU1x1, self).__init__()
        w_shape = [inCh, inCh]
        w_init = torch.qr(torch.randn(*w_shape))[0]

        p, lower, upper = torch.lu_unpack(*torch.lu(w_init))
        s = torch.diag(upper)
        sign_s = torch.sign(s)
        log_s = torch.log(torch.abs(s))
        upper = torch.triu(upper, 1)
        l_mask = torch.tril(torch.ones(w_shape), -1)
        eye = torch.eye(*w_shape)

        self.register_buffer('p', p)  # .cuda() will work only on register_buffer
        self.register_buffer('sign_s', sign_s)
        self.register_buffer('l_mask', l_mask)
        self.register_buffer('eye', eye)

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

        self.w_shape = w_shape
//  [python]
def forward(self, *args, **kwargs):
        # - Get transformed attributes
        transformed_attrs = self._T()

        if self._record and not self._has_torch_api:
            kwargs["record"] = self._record

        # - Call module with torch functional API
        out = torch.nn.utils.stateless.functional_call(
            self._mod, transformed_attrs, args, kwargs
        )

        if not self._has_torch_api:
            self._record_dict = out[2]
        elif hasattr(self._mod, "_record_dict"):
            self._record_dict = self._mod._record_dict

        if not self._has_torch_api:
            return out[0]
        else:
            return out
//  [python]
def forward(self, *args, **kwargs):
        # - Get transformed attributes
        transformed_attrs = self._T()

        if self._record and not self._has_torch_api:
            kwargs["record"] = self._record

        # - Call module with torch functional API
        out = torch.func.functional_call(
            self._mod, transformed_attrs, args, kwargs
        )

        if not self._has_torch_api:
            self._record_dict = out[2]
        elif hasattr(self._mod, "_record_dict"):
            self._record_dict = self._mod._record_dict

        if not self._has_torch_api:
            return out[0]
        else:
            return out
//  [python]
def predict_in_train(self, x):
        if not self._use_lambdas:
            return self.network(x)
        wa_weights = self.get_featurizer_wa_weights()
        features_wa = torch.nn.utils.stateless.functional_call(
                self.featurizer, wa_weights, x)
        return self.classifier(features_wa)
//  [python]
def predict_in_train(self, x):
        if not self._use_lambdas:
            return self.network(x)
        wa_weights = self.get_featurizer_wa_weights()
        features_wa = torch.func.functional_call(
                self.featurizer, wa_weights, x)
        return self.classifier(features_wa)
//  [python]
def train_step(x, y):
    optimizer.zero_grad()

    weights = {}
    for (name_0, param_0), (name_1, param_1) in zip(net_0.named_parameters(), net_1.named_parameters()):
        assert name_0 == name_1
        weights[name_0] = (param_0 * lambda_0 + param_1 * lambda_1) / (lambda_0 + lambda_1)

    preds = torch.nn.utils.stateless.functional_call(net_wa, weights, x)
    loss = loss_fn(preds, y)
    loss.backward(retain_graph=True)
    optimizer.step()
    return loss.item()
//  [python]
def train_step(x, y):
    optimizer.zero_grad()

    weights = {}
    for (name_0, param_0), (name_1, param_1) in zip(net_0.named_parameters(), net_1.named_parameters()):
        assert name_0 == name_1
        weights[name_0] = (param_0 * lambda_0 + param_1 * lambda_1) / (lambda_0 + lambda_1)

    preds = torch.func.functional_call(net_wa, weights, x)
    loss = loss_fn(preds, y)
    loss.backward(retain_graph=True)
    optimizer.step()
    return loss.item()
//  [python]
def predict(self, x, **kwargs):
        dict_predictions = {}

        if kwargs.get("lambdas") is not None:
            lambda_interpolation = [1.] + kwargs.get("lambdas")
        else:
            lambda_interpolation = torch.exp(self.lambdas)

        wa_weights = misc.get_wa_weights(
            lambda_interpolation=lambda_interpolation,
            featurizers=[self.featurizer] + self.featurizers
        )
        features_wa = torch.nn.utils.stateless.functional_call(self.featurizer, wa_weights, x)

        # w for wa, t for task
        dict_predictions["feats"] = features_wa
        dict_predictions["logits"] = self.classifier(features_wa)
        return dict_predictions
//  [python]
def predict(self, x, **kwargs):
        dict_predictions = {}

        if kwargs.get("lambdas") is not None:
            lambda_interpolation = [1.] + kwargs.get("lambdas")
        else:
            lambda_interpolation = torch.exp(self.lambdas)

        wa_weights = misc.get_wa_weights(
            lambda_interpolation=lambda_interpolation,
            featurizers=[self.featurizer] + self.featurizers
        )
        features_wa = torch.func.functional_call(self.featurizer, wa_weights, x)

        # w for wa, t for task
        dict_predictions["feats"] = features_wa
        dict_predictions["logits"] = self.classifier(features_wa)
        return dict_predictions
//  [python]
def __call__(self, inputs, outputs, targets, device):
        if self.method == 'functorch':
            func = lambda params: self.loss_function(
                torch.nn.utils.stateless.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = functorch.jacfwd(functorch.jacrev(func), randomness='same')(self.f_layer)
            return hessian
        elif self.method == 'autograd':
            loss = self.loss_function(outputs, targets)
            shape = self.f_layer.shape
            # need to retain_graph for further hessian calculation
            layer_jacobian = grad(loss, self.f_layer, create_graph=True, retain_graph=True, allow_unused=True)
            drv2 = torch.tensor(torch.empty(shape[1], shape[0], shape[0], shape[1]), requires_grad=True).to(device)
            for ind, n_grd in enumerate(layer_jacobian[0].T):
                for neuron_j in range(shape[0]):
                    drv2[ind][neuron_j] = grad(n_grd[neuron_j].to(device), self.f_layer, retain_graph=True)[0].to(device)
            return drv2
        else:
            print("No such method of hessian computation")
        return None
//  [python]
def __call__(self, inputs, outputs, targets, device):
        if self.method == 'functorch':
            func = lambda params: self.loss_function(
                torch.func.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = functorch.jacfwd(functorch.jacrev(func), randomness='same')(self.f_layer)
            return hessian
        elif self.method == 'autograd':
            loss = self.loss_function(outputs, targets)
            shape = self.f_layer.shape
            # need to retain_graph for further hessian calculation
            layer_jacobian = grad(loss, self.f_layer, create_graph=True, retain_graph=True, allow_unused=True)
            drv2 = torch.tensor(torch.empty(shape[1], shape[0], shape[0], shape[1]), requires_grad=True).to(device)
            for ind, n_grd in enumerate(layer_jacobian[0].T):
                for neuron_j in range(shape[0]):
                    drv2[ind][neuron_j] = grad(n_grd[neuron_j].to(device), self.f_layer, retain_graph=True)[0].to(device)
            return drv2
        else:
            print("No such method of hessian computation")
        return None
//  [python]
def _functional_call(model: Module, d: Dict[str, Tensor], features):
    """
    Makes a call to `model.forward`, which is treated as a function of the parameters
    in `d`, a dict from parameter name to parameter, instead of as a function of
    `features`, the argument that is unpacked to `model.forward` (i.e.
    `model.forward(*features)`).  Depending on what version of PyTorch is available,
    we either use our own implementation, or directly use `torch.nn.utils.stateless`
    or `torch.func.functional_call`.  Put another way, this function mimics the latter
    two implementations, using our own when the PyTorch version is too old.
    """
    import torch

    version = _parse_version(torch.__version__)
    if version < (1, 12, 0):
        return _custom_functional_call(model, d, features)
    elif version >= (1, 12, 0) and version < (2, 0, 0):
        import torch.nn.utils.stateless

        return torch.nn.utils.stateless.functional_call(model, d, features)
    else:
        import torch.func

        return torch.func.functional_call(model, d, features)
//  [python]
def _functional_call(model: Module, d: Dict[str, Tensor], features):
    """
    Makes a call to `model.forward`, which is treated as a function of the parameters
    in `d`, a dict from parameter name to parameter, instead of as a function of
    `features`, the argument that is unpacked to `model.forward` (i.e.
    `model.forward(*features)`).  Depending on what version of PyTorch is available,
    we either use our own implementation, or directly use `torch.nn.utils.stateless`
    or `torch.func.functional_call`.  Put another way, this function mimics the latter
    two implementations, using our own when the PyTorch version is too old.
    """
    import torch

    version = _parse_version(torch.__version__)
    if version < (1, 12, 0):
        return _custom_functional_call(model, d, features)
    else:
        import torch.func

        return torch.func.functional_call(model, d, features)
//  [python]
def test_stateless_functional_call_warns(self):
        m = torch.nn.Linear(1, 1)
        params = dict(m.named_parameters())
        x = torch.randn(3, 1)
        with self.assertWarnsRegex(UserWarning, "Please use torch.func.functional_call"):
            stateless.functional_call(m, params, x)
//  [python]
def test_stateless_functional_call_warns(self):
        m = torch.nn.Linear(1, 1)
        params = dict(m.named_parameters())
        x = torch.randn(3, 1)
        with self.assertWarnsRegex(UserWarning, "Please use torch.func.functional_call"):
            torch.func.functional_call(m, params, x)
//  [python]
def test_reparametrize_tie_some_weights(self, functional_call):
        module = MockTiedModule()
        weight = torch.tensor([[2.0]])
        buffer = torch.tensor([3.0])

        parameters = {'l1.weight': weight,
                      'buffer': buffer}
        x = torch.randn(1, 1)
        out = stateless.functional_call(module, parameters, x, tie_weights=True)
        self.assertEqual(out, x * 2. + module.l1.bias + module.tied_bias + buffer + buffer)
//  [python]
def test_reparametrize_tie_some_weights(self, functional_call):
        module = MockTiedModule()
        weight = torch.tensor([[2.0]])
        buffer = torch.tensor([3.0])

        parameters = {'l1.weight': weight,
                      'buffer': buffer}
        x = torch.randn(1, 1)
        out = torch.func.functional_call(module, parameters, x, tie_weights=True)
        self.assertEqual(out, x * 2. + module.l1.bias + module.tied_bias + buffer + buffer)
//  [python]
def model_to_function(model: torch.nn.Module) -> Callable[[TensorTree, Any], Any]:
    """Converts a model into a function that maps parameters and inputs to outputs.

    Convenience wrapper around [torch.functional_call](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html).

    Args:
        model: torch.nn.Module with parameters stored in .named_parameters().

    Returns:
        Function that takes a PyTree of parameters as well as any input
            arg or kwargs and returns the output of the model.
    """

    def func_model(p_dict, *args, **kwargs):
        return t.functional_call(model, p_dict, args=args, kwargs=kwargs)

    return func_model
//  [python]
def model_to_function(model: torch.nn.Module) -> Callable[[TensorTree, Any], Any]:
    """Converts a model into a function that maps parameters and inputs to outputs.

    Convenience wrapper around [torch.functional_call](https://pytorch.org/docs/stable/generated/torch.func.functional_call.html).

    Args:
        model: torch.nn.Module with parameters stored in .named_parameters().

    Returns:
        Function that takes a PyTree of parameters as well as any input
            arg or kwargs and returns the output of the model.
    """

    def func_model(p_dict, *args, **kwargs):
        return functional_call(model, p_dict, args=args, kwargs=kwargs)

    return func_model
//  [python]
def _functional_call(model, d, features):
    """
    Makes a call to `model.forward`, which is treated as a function of the parameters
    in `d`, a dict from parameter name to parameter, instead of as a function of
    `features`, the argument that is unpacked to `model.forward` (i.e.
    `model.forward(*features)`).  Depending on what version of PyTorch is available,
    we either use our own implementation, or directly use `torch.nn.utils.stateless`
    or `torch.func.functional_call`.  Put another way, this function mimics the latter
    two implementations, using our own when the PyTorch version is too old.
    """
    import torch

    version = torch.__version__
    if version < "1.12":
        return _custom_functional_call(model, d, features)
    elif version >= "1.12" and version < "2.0":
        import torch.nn.utils.stateless

        return torch.nn.utils.stateless.functional_call(model, d, features)
    else:
        import torch.func

        return torch.func.functional_call(model, d, features)
//  [python]
def _functional_call(model, d, features):
    """
    Makes a call to `model.forward`, which is treated as a function of the parameters
    in `d`, a dict from parameter name to parameter, instead of as a function of
    `features`, the argument that is unpacked to `model.forward` (i.e.
    `model.forward(*features)`).  Depending on what version of PyTorch is available,
    we either use our own implementation, or directly use `torch.nn.utils.stateless`
    or `torch.func.functional_call`.  Put another way, this function mimics the latter
    two implementations, using our own when the PyTorch version is too old.
    """
    import torch

    version = torch.__version__
    if version < "1.12":
        return _custom_functional_call(model, d, features)
    else:
        import torch.func

        return torch.func.functional_call(model, d, features)
//  [python]
def __call__(self, inputs, targets, device):
        if self.method == 'functorch':
            func = lambda params: self.loss_function(
                torch.nn.utils.stateless.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = functorch.jacfwd(functorch.jacrev(func), randomness='same')(self.f_layer)
            return hessian
        elif self.method == 'autograd':
            loss = self.loss_function(self.model(inputs), targets)
            shape = self.f_layer.shape
            # need to retain_graph for further hessian calculation
            # need allow_unused for the layer that is transformed from conv2d
            layer_jacobian = grad(loss, self.f_layer, create_graph=True, retain_graph=True, allow_unused=True)
            drv2 = torch.tensor(torch.empty(shape[1], shape[0], shape[0], shape[1]), requires_grad=True).to(device)
            for ind, n_grd in enumerate(layer_jacobian[0].T):
                for neuron_j in range(shape[0]):
                    drv2[ind][neuron_j] = grad(n_grd[neuron_j].to(device), self.f_layer, retain_graph=True)[0].to(device)
            return drv2
        elif self.method == 'autograd_fast':
            func = lambda params: self.loss_function(
                torch.nn.utils.stateless.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = torch.autograd.functional.hessian(func, self.f_layer, vectorize=True, create_graph=True)
            return hessian
        else:
            print("No such method of hessian computation")
        return None
//  [python]
def __call__(self, inputs, targets, device):
        if self.method == 'functorch':
            func = lambda params: self.loss_function(
                torch.func.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = functorch.jacfwd(functorch.jacrev(func), randomness='same')(self.f_layer)
            return hessian
        elif self.method == 'autograd':
            loss = self.loss_function(self.model(inputs), targets)
            shape = self.f_layer.shape
            # need to retain_graph for further hessian calculation
            # need allow_unused for the layer that is transformed from conv2d
            layer_jacobian = grad(loss, self.f_layer, create_graph=True, retain_graph=True, allow_unused=True)
            drv2 = torch.tensor(torch.empty(shape[1], shape[0], shape[0], shape[1]), requires_grad=True).to(device)
            for ind, n_grd in enumerate(layer_jacobian[0].T):
                for neuron_j in range(shape[0]):
                    drv2[ind][neuron_j] = grad(n_grd[neuron_j].to(device), self.f_layer, retain_graph=True)[0].to(device)
            return drv2
        elif self.method == 'autograd_fast':
            func = lambda params: self.loss_function(
                torch.func.functional_call(self.model, {self.f_name: params}, inputs).to(device),
                targets)

            hessian = torch.autograd.functional.hessian(func, self.f_layer, vectorize=True, create_graph=True)
            return hessian
        else:
            print("No such method of hessian computation")
        return None