jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward_dynamicGP(self, Xstar, flg_noise=False):
        """
        inference of dynamicGP, map X0 to X1
        :param Xstar:       X0
        :param flg_noise:   whether add noise
        :return:            X1
        """
        Xin, Xout = self.get_Xin_Xout()
        n = Xstar.shape[0]
        Kx_star = self.dynamicGP_kernel(Xin, Xstar, False)
        mean_Xout_pred = torch.linalg.multi_dot([Xout.t(), self.Kx_inv, Kx_star]).t()
        if flg_noise:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.exp(self.x_log_sigma_n) ** 2 + self.sigma_n_num_X ** 2 + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        else:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        x_log_lambdas = torch.exp(self.x_log_lambdas) ** -2
        diag_var_Xout_pred = diag_var_Xout_pred_common.unsqueeze(1) * x_log_lambdas.unsqueeze(0)

        return mean_Xout_pred, diag_var_Xout_pred
//  [python]
def Qv(self, v: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
        """
        TODO: Check this...

        For fully-connected layers, the default structure of :math:`Q` as a
        block-diaglonal matrix is,
        .. math::
                    Q_l = (R_lR_l^T \otimes L_lL_l^T)
        where :math:`l` denotes the l-th layer. The matrix :math:`R_l` has size
        :math:`(N_{l-1} + 1) \\times (N_{l-1} + 1)` while the matrix :math:`L_l` has
        size :math:`N_l \\times N_l`. The auxiliarary parameters :math:`\lambda`
        are represented by the matrices :math:`L_l, R_l`. For a Kronecker form that
        introduces full inner and outer diagonal rescaling structure is,

        .. math::
                    Q_l = A_l(L_l \otimes R_l^T) D_l^2 (L_l^T \otimes R_l) A_l

        where :math:`A_l` and :math:`D_l` are two diagonal matrices of the
        appropriate size.
        """
        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        u = torch.cat([v[0], v[1][:, None]], dim=-1) if self.bias is not None else v[0]

        A = self.fishleg_aux["A"]
        u = A * u
        u = torch.mm(R.T, torch.mm(R, torch.mm(u, L)))
        u = torch.mm(u,L.T)
        u = A * u
        return (u[:, :-1], u[:, -1]) if self.bias is not None else (u,)
//  [python]
def Qv(self, v: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
        """
        TODO: Check this...

        For fully-connected layers, the default structure of :math:`Q` as a
        block-diaglonal matrix is,
        .. math::
                    Q_l = (R_lR_l^T \otimes L_lL_l^T)
        where :math:`l` denotes the l-th layer. The matrix :math:`R_l` has size
        :math:`(N_{l-1} + 1) \\times (N_{l-1} + 1)` while the matrix :math:`L_l` has
        size :math:`N_l \\times N_l`. The auxiliarary parameters :math:`\lambda`
        are represented by the matrices :math:`L_l, R_l`. For a Kronecker form that
        introduces full inner and outer diagonal rescaling structure is,

        .. math::
                    Q_l = A_l(L_l \otimes R_l^T) D_l^2 (L_l^T \otimes R_l) A_l

        where :math:`A_l` and :math:`D_l` are two diagonal matrices of the
        appropriate size.
        """
        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        u = torch.cat([v[0], v[1][:, None]], dim=-1) if self.bias is not None else v[0]

        A = self.fishleg_aux["A"]
        u = A * u
        u = torch.linalg.multi_dot((R.T, R, u, L, L.T))
        u = A * u
        return (u[:, :-1], u[:, -1]) if self.bias is not None else (u,)
//  [python]
def _precond_grad_dense_dense(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using Kronecker product preconditioner
    Ql: (left side) Cholesky factor of preconditioner
    Qr: (right side) Cholesky factor of preconditioner
    Grad: (matrix) gradient
    """
    return torch.chain_matmul(Ql.t(), Ql, Grad, Qr.t(), Qr)
//  [python]
def _precond_grad_dense_dense(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using Kronecker product preconditioner
    Ql: (left side) Cholesky factor of preconditioner
    Qr: (right side) Cholesky factor of preconditioner
    Grad: (matrix) gradient
    """
    #return torch.chain_matmul(Ql.t(), Ql, Grad, Qr.t(), Qr)
    return torch.linalg.multi_dot([Ql.t(), Ql, Grad, Qr.t(), Qr])
//  [python]
def _complex_python_udf(t1, t2):
        t3 = torch.nn.functional.linear(t1, t2)
        t4 = torch.nn.functional.linear(t2, t3)
        t5 = torch.nn.functional.linear(t3, t4)
        return t1 @ t2 @ t3 @ t4 @ t5
//  [python]
def _python_udf_with_backward_error(t1, t2):
        t3 = t1 + t2
        t4 = SimulateBackwardError.apply(t3)
        return t1.matmul(t2).matmul(t3).matmul(t4)
//  [python]
def _precond_grad_dense_scale(Ql, qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (dense, scaling) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    Ql: shape (M, M), (left side) Cholesky factor of preconditioner
    qr: shape (1, N), defines a diagonal matrix for output feature scaling
    Grad: (matrix) gradient
    """
    return torch.chain_matmul(Ql.t(), Ql, Grad*(qr*qr))
//  [python]
def _precond_grad_dense_scale(Ql, qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (dense, scaling) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    Ql: shape (M, M), (left side) Cholesky factor of preconditioner
    qr: shape (1, N), defines a diagonal matrix for output feature scaling
    Grad: (matrix) gradient
    """
    #return torch.chain_matmul(Ql.t(), Ql, Grad*(qr*qr))
    return torch.linalg.multi_dot([Ql.t(), Ql, Grad*(qr*qr)])
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t()*Grad + ql[1:].t().mm(Grad[-1:]) # Ql*Grad
    preG = torch.chain_matmul(preG, Qr.t(), Qr)
    add_last_row = ql[1:].mm(preG) # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row

    return preG
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t()*Grad + ql[1:].t().mm(Grad[-1:]) # Ql*Grad
    #preG = torch.chain_matmul(preG, Qr.t(), Qr)
    preG = torch.linalg.multi_dot([preG, Qr.t(), Qr])
    add_last_row = ql[1:].mm(preG) # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row

    return preG
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))

    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    Qr *= rho

    A = Ql.mm( dG.mm( Qr.t() ) )
    Bt = torch.triangular_solve((torch.triangular_solve(dX.t(), Qr, upper=True, transpose=True))[0].t(),
                     Ql, upper=True, transpose=True)[0]

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))

    step1 = step/(torch.max(torch.abs(grad1)) + _tiny)
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))

    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    Qr *= rho

    #A = Ql.mm( dG.mm( Qr.t() ) )
    #Bt = torch.triangular_solve((torch.triangular_solve(dX.t(), Qr, upper=True, transpose=True))[0].t(),
    #                 Ql, upper=True, transpose=True)[0]
    A = torch.linalg.multi_dot([Ql, dG, Qr.t()])
    Bt = torch.linalg.solve_triangular(Ql.t(), torch.linalg.solve_triangular(Qr, dX, upper=True, left=False), upper=False)

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))

    step1 = step/(torch.max(torch.abs(grad1)) + _tiny)
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def multidot(H,x,y):
    ret = torch.mm(torch.mm(torch.mm(torch.t(H), y), torch.t(x)), H)
    return ret
//  [python]
def multidot(H,x,y):
    ret = torch.linalg.multi_dot([torch.t(H), y, torch.t(x), H]) 
    return ret
//  [python]
def _nested_rpc_call_backward_error(t1, t2, dst):
        t1 = t1 * t2
        t2 = t1 + t2
        res = rpc.rpc_sync(
            worker_name(dst),
            DistAutogradTest._python_udf_with_backward_error,
            args=(t1, t2),
        )
        return torch.mm(t1, torch.mm(t2, res))
//  [python]
def test_schema_check_mode_functionality_list_input(self):
        a = torch.rand((3, 3))
        b = torch.rand((3, 3))
        c = torch.rand((3, 3))
        expected = torch.mm(torch.mm(a, b), c)
        with SchemaCheckMode():
            actual = torch.mm(torch.mm(a, b), c)
        self.assertEqual(expected, actual)
//  [python]
def test_schema_check_mode_functionality_list_input(self):
        a = torch.rand((3, 3))
        b = torch.rand((3, 3))
        c = torch.rand((3, 3))
        expected = torch.linalg.multi_dot([a, b, c])
        with SchemaCheckMode():
            actual = torch.linalg.multi_dot([a, b, c])
        self.assertEqual(expected, actual)
//  [python]
def multi_dot(inputs: Sequence[ArrayLike], *, out=None):
    return t.linalg.multi_dot(inputs)
//  [python]
def multi_dot(inputs: Sequence[ArrayLike], *, out=None):
    return torch.linalg.multi_dot(inputs)
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes):
            tensors = [make_tensor(shape, device, dtype) for shape in shapes]
            np_arrays = [tensor.cpu().numpy() for tensor in tensors]
            res = t.mm(tensors[0], t.mm(tensors[1], t.mm(tensors[2],t.mm(tensors[3], tensors[4])))) if len(tensors) == 5 else t.mm(tensors[0], t.mm(tensors[1], tensors[2])) if len(tensors) == 3 else t.mm(tensors[0], tensors[1]) if len(tensors) == 2 else tensors[0]
            ref = torch.from_numpy(np.array(np.linalg.multi_dot(np_arrays)))
            self.assertEqual(res, ref)

        # test for inputs with empty dimensions
        check([0], [0])
        check([2], [2, 0])
        check([1, 0], [0])
        check([0, 2], [2, 1])
        check([2, 2], [2, 0])
        check([2, 0], [0, 3])
        check([0, 0], [0, 1])
        check([4, 2], [2, 0], [0, 3], [3, 2])

        # test variable output shapes
        check([2], [2])
        check([1, 2], [2])
        check([2], [2, 1])
        check([1, 2], [2, 1])
        check([3, 2], [2, 4])

        # test multiple input tensors
        check([3], [3, 4], [4, 2], [2, 5], [5])
        check([1, 2], [2, 2], [2, 3], [3, 1])

        # test large tensors
        check([10, 100], [100, 5], [5, 50])
        check([10, 20], [20, 30], [30, 5])
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes):
            tensors = [make_tensor(shape, device, dtype) for shape in shapes]
            np_arrays = [tensor.cpu().numpy() for tensor in tensors]
            res = torch.linalg.multi_dot(tensors).cpu()
            ref = torch.from_numpy(np.array(np.linalg.multi_dot(np_arrays)))
            self.assertEqual(res, ref)

        # test for inputs with empty dimensions
        check([0], [0])
        check([2], [2, 0])
        check([1, 0], [0])
        check([0, 2], [2, 1])
        check([2, 2], [2, 0])
        check([2, 0], [0, 3])
        check([0, 0], [0, 1])
        check([4, 2], [2, 0], [0, 3], [3, 2])

        # test variable output shapes
        check([2], [2])
        check([1, 2], [2])
        check([2], [2, 1])
        check([1, 2], [2, 1])
        check([3, 2], [2, 4])

        # test multiple input tensors
        check([3], [3, 4], [4, 2], [2, 5], [5])
        check([1, 2], [2, 2], [2, 3], [3, 1])

        # test large tensors
        check([10, 100], [100, 5], [5, 50])
        check([10, 20], [20, 30], [30, 5])
//  [python]
def get_y_neg_log_likelihood(self, Y, X, N):
        """
        L_y = D/2*log(|K_y(X,X)|) + 1/2*trace(K_y^-1*Y*W_y^2*Y) - N*log(|W_y|)
        """
        K_y = self.observationGP_kernel(X, X)
        U, info = torch.linalg.cholesky_ex(K_y, upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        log_det_K_y = 2 * torch.sum(torch.log(torch.diag(U)))

        W2 = torch.diag(torch.exp(self.y_log_lambdas) ** 2)
        log_det_W = 2 * torch.sum(self.y_log_lambdas)

        return self.D / 2 * log_det_K_y + 1 / 2 * \
               torch.trace(torch.mm(torch.mm(torch.mm(Ky_inv, Y), W2), Y.t())) \
               - N * log_det_W
//  [python]
def get_y_neg_log_likelihood(self, Y, X, N):
        """
        L_y = D/2*log(|K_y(X,X)|) + 1/2*trace(K_y^-1*Y*W_y^2*Y) - N*log(|W_y|)
        """
        K_y = self.observationGP_kernel(X, X)
        U, info = torch.linalg.cholesky_ex(K_y, upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        log_det_K_y = 2 * torch.sum(torch.log(torch.diag(U)))

        W2 = torch.diag(torch.exp(self.y_log_lambdas) ** 2)
        log_det_W = 2 * torch.sum(self.y_log_lambdas)

        return self.D / 2 * log_det_K_y + 1 / 2 * \
               torch.trace(torch.linalg.multi_dot([Ky_inv, Y, W2, Y.t()])) \
               - N * log_det_W
//  [python]
def Qv(self, v: Tuple[Tensor,], full: bool = False) -> Tuple[Tensor,]:
        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        u = v[0]

        A = self.fishleg_aux["A"]
        u = A * u
        u = torch.matmul(torch.diag(R), torch.matmul(torch.diag(R), torch.matmul(u, torch.matmul(L, L.T))))
        u = A * u
        return (u,)
//  [python]
def Qv(self, v: Tuple[Tensor,], full: bool = False) -> Tuple[Tensor,]:
        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        u = v[0]

        A = self.fishleg_aux["A"]
        u = A * u
        u = torch.linalg.multi_dot((torch.diag(R), torch.diag(R), u, L, L.T))
        u = A * u
        return (u,)
//  [python]
def calculate_inverse(matrix, epsilon=1e-6, return_sqrt=False, mode='regularize'):
    """ This method can be applied to compute the inverse or the square-root of the inverse of the matrix,
    this method will be further used to estimate koopman matrix.

    Parameters
    ----------
    matrix : torch.Tensor
        The matrix to be inverted.
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    return_sqrt : boolean, optional, default = False
        If True, the square root of the inverse matrix is returned instead.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    inverse : torch.Tensor
        Inverse of the matrix.
    """

    eigval, eigvec = eig_decomposition(matrix, epsilon, mode)

    if return_sqrt:
        diag = torch.diag(torch.sqrt(1. / eigval))
    else:
        diag = torch.diag(1. / eigval)

    inverse = torch.chain_matmul(eigvec.t(), diag, eigvec)

    return inverse
//  [python]
def calculate_inverse(matrix, epsilon=1e-6, return_sqrt=False, mode='regularize'):
    """ This method can be applied to compute the inverse or the square-root of the inverse of the matrix,
    this method will be further used to estimate koopman matrix.

    Parameters
    ----------
    matrix : torch.Tensor
        The matrix to be inverted.
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    return_sqrt : boolean, optional, default = False
        If True, the square root of the inverse matrix is returned instead.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    inverse : torch.Tensor
        Inverse of the matrix.
    """

    eigval, eigvec = eig_decomposition(matrix, epsilon, mode)

    if return_sqrt:
        diag = torch.diag(torch.sqrt(1. / eigval))
    else:
        diag = torch.diag(1. / eigval)

    try:
    # inverse = torch.chain_matmul(eigvec.t(), diag, eigvec)
        inverse = torch.linalg.multi_dot((eigvec.t(), diag, eigvec))
    except:
        inverse = torch.chain_matmul(eigvec.t(), diag, eigvec)

    return inverse
//  [python]
def forward_dynamicGP(self, Xstar, flg_noise=False):
        Xin, Xout = self.get_Xin_Xout()
        n = Xstar.shape[0]
        Kx_star = self.dynamicGP_kernel(Xin, Xstar, False)
        mean_Xout_pred = (Xout.t() @ self.Kx_inv @ Kx_star).t()
        if flg_noise:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.exp(self.x_log_sigma_n) ** 2 + self.sigma_n_num_X ** 2 + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        else:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        x_log_lambdas = torch.exp(self.x_log_lambdas) ** -2
        diag_var_Xout_pred = diag_var_Xout_pred_common.unsqueeze(1) * x_log_lambdas.unsqueeze(0)

        return mean_Xout_pred, diag_var_Xout_pred
//  [python]
def forward_dynamicGP(self, Xstar, flg_noise=False):
        Xin, Xout = self.get_Xin_Xout()
        n = Xstar.shape[0]
        Kx_star = self.dynamicGP_kernel(Xin, Xstar, False)
        mean_Xout_pred = torch.linalg.multi_dot([Xout.t(), self.Kx_inv, Kx_star]).t()
        if flg_noise:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.exp(self.x_log_sigma_n) ** 2 + self.sigma_n_num_X ** 2 + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        else:
            diag_var_Xout_pred_common = torch.ones(n, dtype=self.dtype, device=self.device) + \
                                        torch.sum(self.log_coe_linear.exp() * Xstar * Xstar) - \
                                        torch.sum(Kx_star.t() @ self.Kx_inv * Kx_star.t(), dim=1)
        x_log_lambdas = torch.exp(self.x_log_lambdas) ** -2
        diag_var_Xout_pred = diag_var_Xout_pred_common.unsqueeze(1) * x_log_lambdas.unsqueeze(0)

        return mean_Xout_pred, diag_var_Xout_pred
//  [python]
def forward_observationGP(self, Xstar, flg_noise=False):
        U, _ = torch.linalg.cholesky_ex(self.observationGP_kernel(self.X, self.X), upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        Y_obs = torch.tensor(np.concatenate(self.observations_list, 0), dtype=self.dtype,
                             device=self.device)  # (M*N, D)
        Ky_star = self.observationGP_kernel(self.X, Xstar, False)
        mean_Y_pred = (Y_obs.t() @ Ky_inv @ Ky_star).t()
        if flg_noise:
            diag_var_Y_pred_common = torch.ones(Xstar.shape[0], dtype=self.dtype, device=self.device) + \
                                     torch.exp(self.y_log_sigma_n) ** 2 + self.sigma_n_num_Y ** 2
        else:
            diag_var_Y_pred_common = torch.ones(Xstar.shape[0], dtype=self.dtype, device=self.device)
        diag_var_Y_pred_common = diag_var_Y_pred_common - torch.sum(Ky_star.t() @ Ky_inv * Ky_star.t(),
                                                                    dim=1)
        y_log_lambdas = torch.exp(self.y_log_lambdas) ** -2
        diag_var_Y_pred = diag_var_Y_pred_common.unsqueeze(1) * y_log_lambdas.unsqueeze(0)

        return mean_Y_pred, diag_var_Y_pred
//  [python]
def forward_observationGP(self, Xstar, flg_noise=False):
        U, _ = torch.linalg.cholesky_ex(self.observationGP_kernel(self.X, self.X), upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        Y_obs = torch.tensor(np.concatenate(self.observations_list, 0), dtype=self.dtype,
                             device=self.device)  # (M*N, D)
        Ky_star = self.observationGP_kernel(self.X, Xstar, False)
        mean_Y_pred = torch.linalg.multi_dot([Y_obs.t(), Ky_inv, Ky_star]).t()
        if flg_noise:
            diag_var_Y_pred_common = torch.ones(Xstar.shape[0], dtype=self.dtype, device=self.device) + \
                                     torch.exp(self.y_log_sigma_n) ** 2 + self.sigma_n_num_Y ** 2
        else:
            diag_var_Y_pred_common = torch.ones(Xstar.shape[0], dtype=self.dtype, device=self.device)
        diag_var_Y_pred_common = diag_var_Y_pred_common - torch.sum(Ky_star.t() @ Ky_inv * Ky_star.t(),
                                                                    dim=1)
        y_log_lambdas = torch.exp(self.y_log_lambdas) ** -2
        diag_var_Y_pred = diag_var_Y_pred_common.unsqueeze(1) * y_log_lambdas.unsqueeze(0)

        return mean_Y_pred, diag_var_Y_pred
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, flg_W=True):
        """
        If use scaling matrix W, choose flg_W = True
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - (N-dyn_back_step)*log(|W_x|)

        If use original GPDM, choose flg_W = False
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*Xout)

        ISSUE: W_x is necessary
        """

        K_x = self.dynamicGP_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        if flg_W:
            W2 = torch.diag(torch.exp(self.x_log_lambdas) ** 2)
            log_det_W = 2 * torch.sum(self.x_log_lambdas)
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.matmul(torch.matmul(Kx_inv, Xout), torch.matmul(W2, Xout.t()))) \
                   - Xin.shape[0] * log_det_W
        else:
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.matmul(torch.matmul(Kx_inv, Xout), Xout.t()))
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, flg_W=True):
        """
        If use scaling matrix W, choose flg_W = True
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - (N-dyn_back_step)*log(|W_x|)

        If use original GPDM, choose flg_W = False
        L_x = d/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*Xout)

        ISSUE: W_x is necessary
        """

        K_x = self.dynamicGP_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        if flg_W:
            W2 = torch.diag(torch.exp(self.x_log_lambdas) ** 2)
            log_det_W = 2 * torch.sum(self.x_log_lambdas)
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.linalg.multi_dot([Kx_inv, Xout, W2, Xout.t()])) \
                   - Xin.shape[0] * log_det_W
        else:
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.linalg.multi_dot([Kx_inv, Xout, Xout.t()]))
//  [python]
def guess_hess_q(
        self, coords_x: torch.Tensor, atomic_nums: torch.Tensor
    ) -> torch.Tensor:
        import tico.hess

        hess_q = tico.hess.guess_hess_q(coords_x, self.idxs, atomic_nums)
        return torch.mm(self.v.T, torch.mm(hess_q, self.v))
//  [python]
def guess_hess_q(
        self, coords_x: torch.Tensor, atomic_nums: torch.Tensor
    ) -> torch.Tensor:
        import tico.hess

        hess_q = tico.hess.guess_hess_q(coords_x, self.idxs, atomic_nums)
        return torch.linalg.multi_dot([self.v.T, hess_q, self.v])
//  [python]
def __call__(self, coef):
        """
        Evaluate the expansion

        Parameters
        ----------
        coef : np.ndarray (nterms, nsamples)
            The coefficients of the KLE basis
        """
        assert coef.ndim == 2
        assert coef.shape[0] == self.nterms
        if self.use_log:
            if not self.use_torch:
                return np.exp(self.mean_field[:, None]+self.eig_vecs.dot(coef))
            import torch
            return torch.exp(
                self.mean_field[:, None] +
                torch.mm(torch.as_tensor(self.eig_vecs), coef))
        if not self.use_torch:
            return self.mean_field[:, None] + self.eig_vecs.dot(coef)
        import torch
        return (self.mean_field[:, None] +
                torch.mm(torch.as_tensor(self.eig_vecs), coef))
//  [python]
def __call__(self, coef):
        """
        Evaluate the expansion

        Parameters
        ----------
        coef : np.ndarray (nterms, nsamples)
            The coefficients of the KLE basis
        """
        assert coef.ndim == 2
        assert coef.shape[0] == self.nterms
        if self.use_log:
            if not self.use_torch:
                return np.exp(self.mean_field[:, None]+self.eig_vecs.dot(coef))
            import torch
            return torch.exp(
                self.mean_field[:, None] +
                torch.linalg.multi_dot((torch.as_tensor(self.eig_vecs), coef)))
        if not self.use_torch:
            return self.mean_field[:, None] + self.eig_vecs.dot(coef)
        import torch
        return (self.mean_field[:, None] +
                torch.linalg.multi_dot((torch.as_tensor(self.eig_vecs), coef)))
//  [python]
def estimate_koopman_matrix(data: torch.Tensor, data_lagged: torch.Tensor, epsilon=1e-6, mode='regularize', symmetrized=False):
    """ This method can be applied to compute the koopman matrix from time-instant and time-lagged data.

    Parameters
    ----------
    data : torch.Tensor
        The time-instant data of shape [batch_size, num_basis].
    data_lagged : torch.Tensor
        The time-lagged data of shape [batch_size, num_basis].
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    koopman_matrix : torch.Tensor
        The koopman matrix of shape [num_basis, num_basis].
    """

    cov_00, cov_01, cov_11 = compute_covariance_matrix(data, data_lagged)

    if not symmetrized:
        cov_00_sqrt_inverse = calculate_inverse(cov_00, epsilon=epsilon, return_sqrt=True, mode=mode)
        cov_11_sqrt_inverse = calculate_inverse(cov_11, epsilon=epsilon, return_sqrt=True, mode=mode)
        koopman_matrix = torch.chain_matmul(cov_00_sqrt_inverse, cov_01, cov_11_sqrt_inverse).t()
    else:
        cov_0 = 0.5*(cov_00+cov_11)
        cov_1 = 0.5*(cov_01+cov_01.t())
        cov_0_sqrt_inverse = calculate_inverse(cov_0, epsilon=epsilon, return_sqrt=True, mode=mode)
        koopman_matrix = torch.chain_matmul((cov_0_sqrt_inverse, cov_1, cov_0_sqrt_inverse)).t()

    return koopman_matrix
//  [python]
def estimate_koopman_matrix(data: torch.Tensor, data_lagged: torch.Tensor, epsilon=1e-6, mode='regularize', symmetrized=False):
    """ This method can be applied to compute the koopman matrix from time-instant and time-lagged data.

    Parameters
    ----------
    data : torch.Tensor
        The time-instant data of shape [batch_size, num_basis].
    data_lagged : torch.Tensor
        The time-lagged data of shape [batch_size, num_basis].
    epsilon : float, default = 1e-6
        The regularization/trunction parameters for eigenvalues.
    mode : str, default = 'regularize'
        'regularize': regularize the eigenvalues by adding epsilon.
        'trunc': truncate the eigenvalues by filtering out the eigenvalues below epsilon.

    Returns
    -------
    koopman_matrix : torch.Tensor
        The koopman matrix of shape [num_basis, num_basis].
    """

    cov_00, cov_01, cov_11 = compute_covariance_matrix(data, data_lagged)

    if not symmetrized:
        cov_00_sqrt_inverse = calculate_inverse(cov_00, epsilon=epsilon, return_sqrt=True, mode=mode)
        cov_11_sqrt_inverse = calculate_inverse(cov_11, epsilon=epsilon, return_sqrt=True, mode=mode)
        try:
            koopman_matrix = torch.linalg.multi_dot((cov_00_sqrt_inverse, cov_01, cov_11_sqrt_inverse)).t()
        except:
            koopman_matrix = torch.chain_matmul(cov_00_sqrt_inverse, cov_01, cov_11_sqrt_inverse).t()
    else:
        cov_0 = 0.5*(cov_00+cov_11)
        cov_1 = 0.5*(cov_01+cov_01.t())
        cov_0_sqrt_inverse = calculate_inverse(cov_0, epsilon=epsilon, return_sqrt=True, mode=mode)
        try:
            koopman_matrix = torch.linalg.multi_dot((cov_0_sqrt_inverse, cov_1, cov_0_sqrt_inverse)).t()
        except:
            koopman_matrix = torch.chain_matmul((cov_0_sqrt_inverse, cov_1, cov_0_sqrt_inverse)).t()

    return koopman_matrix
//  [python]
def _nested_python_udf(t1, t2, dst):
        t3 = t1 * t2
        t4 = t1 + t2
        res = rpc.rpc_sync(worker_name(dst), my_py_add, args=(t3, t4))
        return t.mm(t.mm(t.mm(t1, t2), t3), t.mm(t4, res))
//  [python]
def project_grad_x(
        self, coords_x: torch.Tensor, grad_x: torch.Tensor
    ) -> torch.Tensor:
        """Project out the components of the gradient along the constrained degrees of
        freedom."""

        if self.constr is None:
            return grad_x

        b_matrix = self.compute_b(coords_x)
        g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

        n_constr = sum(len(v) for i, v in self.constr.values())

        grad_q = t.mm(t.mm(g_matrix_inv, b_matrix), grad_x)
        grad_q[:n_constr] = 0.0

        return t.mm(b_matrix.T, grad_q)
//  [python]
def project_grad_x(
        self, coords_x: torch.Tensor, grad_x: torch.Tensor
    ) -> torch.Tensor:
        """Project out the components of the gradient along the constrained degrees of
        freedom."""

        if self.constr is None:
            return grad_x

        b_matrix = self.compute_b(coords_x)
        g_matrix_inv = tico.utils.pinv(b_matrix @ b_matrix.T)

        n_constr = sum(len(v) for i, v in self.constr.values())

        grad_q = torch.linalg.multi_dot([g_matrix_inv, b_matrix, grad_x])
        grad_q[:n_constr] = 0.0

        return torch.linalg.multi_dot([b_matrix.T, grad_q])
//  [python]
def _precond_grad_dense_dense(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using Kronecker product preconditioner
    Ql: (left side) Cholesky factor of preconditioner
    Qr: (right side) Cholesky factor of preconditioner
    Grad: (matrix) gradient
    """
    return Ql.t() @ Ql @ Grad @ Qr.t() @ Qr
//  [python]
def _precond_grad_dense_dense(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using Kronecker product preconditioner
    Ql: (left side) Cholesky factor of preconditioner
    Qr: (right side) Cholesky factor of preconditioner
    Grad: (matrix) gradient
    """
    return torch.linalg.multi_dot([Ql.t(), Ql, Grad, Qr.t(), Qr])
//  [python]
def _update(self, s, y, rho_inv):
        rho = rho_inv.reciprocal()
        if self.inverse:
            if self.n_updates == 0:
                self.H.mul_(rho_inv / y.dot(y))
            R = torch.addr(self.I, s, y, alpha=-rho)
            torch.addr(
                torch.mm(torch.mm(R, self.H), R.t()),
                s, s, alpha=rho, out=self.H)
        else:
            if self.n_updates == 0:
                self.B.mul_(rho * y.dot(y))
            Bs = torch.mv(self.B, s)
            self.B.addr_(y, y, alpha=rho)
            self.B.addr_(Bs, Bs, alpha=-1./s.dot(Bs))
//  [python]
def _update(self, s, y, rho_inv):
        rho = rho_inv.reciprocal()
        if self.inverse:
            if self.n_updates == 0:
                self.H.mul_(rho_inv / y.dot(y))
            R = torch.addr(self.I, s, y, alpha=-rho)
            torch.addr(
                torch.linalg.multi_dot((R, self.H, R.t())),
                s, s, alpha=rho, out=self.H)
        else:
            if self.n_updates == 0:
                self.B.mul_(rho * y.dot(y))
            Bs = torch.mv(self.B, s)
            self.B.addr_(y, y, alpha=rho)
            self.B.addr_(Bs, Bs, alpha=-1./s.dot(Bs))
//  [python]
def test_backwards_nested_python_udf(self):
        # Run equivalent of _nested_python_udf locally.
        t1 = torch.rand((3, 3), requires_grad=True)
        t2 = torch.rand((3, 3), requires_grad=True)
        t3 = t1 * t2
        t4 = t1 + t2
        res = t3 + t4
        loss = torch.mm(t1, torch.mm(t2, torch.mm(t3, torch.mm(t4, res)))).sum()
        torch.autograd.backward([loss])

        # Now run distributed autograd.
        with dist_autograd.context() as context_id:
            loss = rpc.rpc_sync(
                worker_name(self._next_rank()),
                DistAutogradTest._nested_python_udf,
                args=(t1, t2, self._next_rank()),
            )
            dist_autograd.backward(context_id, [loss.sum()])

            grads = dist_autograd.get_gradients(context_id)
            self.assertEqual(t1.grad, grads[t1])
            self.assertEqual(t2.grad, grads[t2])
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, flg_W=True):
        """
        If use scaling matrix W, choose flg_W = True
        L_x = Q/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - N*log(|W_x|)

        If use original GPDM, choose flg_W = False
        L_x = Q/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*Xout)

        ISSUE: W_x is necessary
        """

        K_x = self.dynamicGP_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        if flg_W:
            W2 = torch.diag(torch.exp(self.x_log_lambdas) ** 2)
            log_det_W = 2 * torch.sum(self.x_log_lambdas)
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(Kx_inv @ Xout @ W2 @ Xout.t()) \
                   - Xin.shape[0] * log_det_W
        else:
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(Kx_inv @ Xout @ Xout.t())
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin, flg_W=True):
        """
        If use scaling matrix W, choose flg_W = True
        L_x = Q/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*W_x^2*Xout) - N*log(|W_x|)

        If use original GPDM, choose flg_W = False
        L_x = Q/2*log(|K_x(Xin,Xin)|) + 1/2*trace(K_x^-1*Xout*Xout)

        ISSUE: W_x is necessary
        """

        K_x = self.dynamicGP_kernel(Xin, Xin)
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        log_det_K_x = 2 * torch.sum(torch.log(torch.diag(U)))

        if flg_W:
            W2 = torch.diag(torch.exp(self.x_log_lambdas) ** 2)
            log_det_W = 2 * torch.sum(self.x_log_lambdas)
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.linalg.multi_dot([Kx_inv, Xout, W2, Xout.t()])) \
                   - Xin.shape[0] * log_det_W
        else:
            return self.Q / 2 * log_det_K_x + 1 / 2 * \
                   torch.trace(torch.linalg.multi_dot([Kx_inv, Xout, Xout.t()]))
//  [python]
def mat(self):
        v_outer = torch.einsum('fi,fj->fij', self.v, self.v)
        v_norms_squared = torch.linalg.norm(self.v, dim=1).view(-1, 1, 1) ** 2
        h = (torch.eye(self.n_dim)[None] - 2 * (v_outer / v_norms_squared))
        return torch.mm(*h)
//  [python]
def mat(self):
        v_outer = torch.einsum('fi,fj->fij', self.v, self.v)
        v_norms_squared = torch.linalg.norm(self.v, dim=1).view(-1, 1, 1) ** 2
        h = (torch.eye(self.n_dim)[None] - 2 * (v_outer / v_norms_squared))
        return torch.linalg.multi_dot(list(h))
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner 
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t()*Grad + ql[1:].t().mm(Grad[-1:]) # Ql*Grad 
    preG = torch.chain_matmul(preG, Qr.t(), Qr)
    add_last_row = ql[1:].mm(preG) # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row
    
    return preG
//  [python]
def _precond_grad_norm_dense(ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    return preconditioned gradient using (normalization, dense) Kronecker product preconditioner 
    Suppose Grad has shape (M, N)
    ql[0] is the diagonal part of Ql
    ql[1, 0:-1] is the last column of Ql, excluding the last entry
    Qr: shape (N, N), Cholesky factor of right preconditioner
    Grad: (matrix) gradient
    """
    preG = ql[0:1].t()*Grad + ql[1:].t().mm(Grad[-1:]) # Ql*Grad 
    #preG = torch.chain_matmul(preG, Qr.t(), Qr)
    preG = torch.linalg.multi_dot([preG, Qr.t(), Qr])
    add_last_row = ql[1:].mm(preG) # use it to modify the last row
    preG *= ql[0:1].t()
    preG[-1:] += add_last_row
    
    return preG
//  [python]
def precond_grad_affine_math(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    Basically a copy of function 
        _precond_grad_dense_dense(Ql, Qr, Grad)
    but with support of complex matrices and diagonal matrices 
    """
    if Ql.dim()==2:
        if Qr.dim()==2: # Ql.dim()=2 and Qr.dim()=2:
            return torch.mm(torch.mm(Ql.H, Ql), torch.mm(Grad, torch.mm(Qr.H, Qr)))
        else: # Ql.dim()=2 and Qr.dim()=1:
            return torch.mm(torch.mm(Ql.H, Ql), Grad*(Qr*Qr.conj()))
    else:
        if Qr.dim()==2: # Ql.dim()=1 and Qr.dim()=2:
            return torch.mm((Ql*Ql.conj())[:,None] * Grad, torch.mm(Qr.H, Qr))
        else: # Ql.dim()=1 and Qr.dim()=1:
            return (Ql*Ql.conj())[:,None] * Grad * (Qr*Qr.conj())
//  [python]
def precond_grad_affine_math(Ql, Qr, Grad):
    # type: (Tensor, Tensor, Tensor) -> Tensor
    """
    Basically a copy of function 
        _precond_grad_dense_dense(Ql, Qr, Grad)
    but with support of complex matrices and diagonal matrices 
    """
    if Ql.dim()==2:
        if Qr.dim()==2: # Ql.dim()=2 and Qr.dim()=2:
            return torch.linalg.multi_dot([Ql.H, Ql, Grad, Qr.H, Qr])
        else: # Ql.dim()=2 and Qr.dim()=1:
            return torch.linalg.multi_dot([Ql.H, Ql, Grad*(Qr*Qr.conj())])
    else:
        if Qr.dim()==2: # Ql.dim()=1 and Qr.dim()=2:
            return torch.linalg.multi_dot([(Ql*Ql.conj())[:,None] * Grad, Qr.H, Qr])
        else: # Ql.dim()=1 and Qr.dim()=1:
            return (Ql*Ql.conj())[:,None] * Grad * (Qr*Qr.conj())
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1] 
    _tiny: an offset to avoid division by zero 
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    Qr *= rho
    
    A = Ql.mm(dG.mm(Qr.t()))
    Bt = torch.triangular_solve((torch.triangular_solve(dX.t(), Qr, upper=True, transpose=True))[0].t(), 
                     Ql, upper=True, transpose=True)[0]
    
    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    
    step1 = step/(norm_lower_bound(grad1) + _tiny)
    step2 = step/(norm_lower_bound(grad2) + _tiny)
        
    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1] 
    _tiny: an offset to avoid division by zero 
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    Qr *= rho
    
    #A = Ql.mm( dG.mm( Qr.t() ) )
    #Bt = torch.triangular_solve((torch.triangular_solve(dX.t(), Qr, upper=True, transpose=True))[0].t(), 
    #                 Ql, upper=True, transpose=True)[0]
    A = torch.linalg.multi_dot([Ql, dG, Qr.t()])
    Bt = torch.linalg.solve_triangular(Ql.t(), torch.linalg.solve_triangular(Qr, dX, upper=True, left=False), upper=False)
    
    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    
    # step1 = step/(torch.max(torch.abs(grad1)) + _tiny)
    # step2 = step/(torch.max(torch.abs(grad2)) + _tiny)
    step1 = step/(norm_lower_bound(grad1) + _tiny)
    step2 = step/(norm_lower_bound(grad2) + _tiny)
        
    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes):
            tensors = [make_tensor(shape, dtype=dtype, device=device) for shape in shapes]
            np_arrays = [tensor.cpu().numpy() for tensor in tensors]
            res = t.mm(tensors[0], t.mm(tensors[1], t.mm(tensors[2], t.mm(tensors[3], tensors[4])))) if len(tensors) == 5 else (t.mm(tensors[0], t.mm(tensors[1], tensors[2])) if len(tensors) == 3 else (t.mm(tensors[0], t.mm(tensors[1], tensors[2])) if len(tensors) == 3 else (t.mm(tensors[0], tensors[1]) if len(tensors) == 2 else tensors[0])))
            ref = torch.from_numpy(np.array(np.linalg.multi_dot(np_arrays)))
            self.assertEqual(res, ref)

        # test for inputs with empty dimensions
        check([0], [0])
        check([2], [2, 0])
        check([1, 0], [0])
        check([0, 2], [2, 1])
        check([2, 2], [2, 0])
        check([2, 0], [0, 3])
        check([0, 0], [0, 1])
        check([4, 2], [2, 0], [0, 3], [3, 2])

        # test variable output shapes
        check([2], [2])
        check([1, 2], [2])
        check([2], [2, 1])
        check([1, 2], [2, 1])
        check([3, 2], [2, 4])

        # test multiple input tensors
        check([3], [3, 4], [4, 2], [2, 5], [5])
        check([1, 2], [2, 2], [2, 3], [3, 1])

        # test large tensors
        check([10, 100], [100, 5], [5, 50])
        check([10, 20], [20, 30], [30, 5])
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes):
            tensors = [make_tensor(shape, dtype=dtype, device=device) for shape in shapes]
            np_arrays = [tensor.cpu().numpy() for tensor in tensors]
            res = torch.linalg.multi_dot(tensors).cpu()
            ref = torch.from_numpy(np.array(np.linalg.multi_dot(np_arrays)))
            self.assertEqual(res, ref)

        # test for inputs with empty dimensions
        check([0], [0])
        check([2], [2, 0])
        check([1, 0], [0])
        check([0, 2], [2, 1])
        check([2, 2], [2, 0])
        check([2, 0], [0, 3])
        check([0, 0], [0, 1])
        check([4, 2], [2, 0], [0, 3], [3, 2])

        # test variable output shapes
        check([2], [2])
        check([1, 2], [2])
        check([2], [2, 1])
        check([1, 2], [2, 1])
        check([3, 2], [2, 4])

        # test multiple input tensors
        check([3], [3, 4], [4, 2], [2, 5], [5])
        check([1, 2], [2, 2], [2, 3], [3, 1])

        # test large tensors
        check([10, 100], [100, 5], [5, 50])
        check([10, 20], [20, 30], [30, 5])
//  [python]
def multidot(arrays):
    return t.linalg.multi_dot(arrays)
//  [python]
def multidot(arrays):
    return torch.linalg.multi_dot(arrays)
//  [python]
def test_shape_inference_unsupported_symbols(self):
        """Test Glow shape inference unsupported symbols."""

        def f(a):
            # linalg.multi_dot is currently not supported by shape inference engine
            return torch.matrix_power(torch.mm(a * 3, a + 4), 3)

        a = torch.randn(3, 3)

        jit_f = torch.jit.trace(f, (a))
        jit_f_graph = jit_f.graph_for(a)

        args = (a,)

        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args
        )
        expected = ["aten::mm", "aten::linalg_matrix_power"]
        self.assertEqual(set(expected), set(actual))

        blocklist = ["aten::mm"]
        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args, blocklist
        )
        expected = ["aten::linalg_matrix_power"]
        self.assertEqual(set(expected), set(actual))
//  [python]
def test_shape_inference_unsupported_symbols(self):
        """Test Glow shape inference unsupported symbols."""

        def f(a):
            # linalg.multi_dot is currently not supported by shape inference engine
            return torch.matrix_power(torch.linalg.multi_dot([a * 3, a + 4]), 3)

        a = torch.randn(3, 3)

        jit_f = torch.jit.trace(f, (a))
        jit_f_graph = jit_f.graph_for(a)

        args = (a,)

        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args
        )
        expected = ["aten::linalg_multi_dot", "aten::linalg_matrix_power"]
        self.assertEqual(set(expected), set(actual))

        blocklist = ["aten::linalg_multi_dot"]
        actual = torch_glow.glow_shape_inference_find_unsupported_symbols(
            jit_f_graph, args, blocklist
        )
        expected = ["aten::linalg_matrix_power"]
        self.assertEqual(set(expected), set(actual))