jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def Qv(
        self, v: Tuple[Tensor, Optional[Tensor]], full: bool = False
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Inspired by KFAC's conv2D layer by Grosse and Martens: Kronecker product of sizes (out_channels ⊗  (in_channels_eff * k_size))
        """

        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        A = self.fishleg_aux["A"]

        if self._bias:
            w, b = v
            sw = w.shape
            sb = b.shape
            w = torch.reshape(w, (self.out_channels, -1))
            b = torch.reshape(b, (self.out_channels, 1))
            u = torch.cat([w, b], dim=-1)
        else:
            (w,) = v
            sw = w.shape
            u = torch.reshape(w, (self.out_channels, -1))

        # at this stage, u is out_channels * (in_channels_eff * k_size (perhaps +1))

        u = torch.matmul(R, torch.matmul(A * u, L.T))
        u = A * torch.matmul(R.T, torch.matmul(u, L))

        if self._bias:
            return (torch.reshape(u[:, :-1], sw), torch.reshape(u[:, -1], sb))
        else:
            return (torch.reshape(u, sw),)
//  [python]
def Qv(
        self, v: Tuple[Tensor, Optional[Tensor]], full: bool = False
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """
        Inspired by KFAC's conv2D layer by Grosse and Martens: Kronecker product of sizes (out_channels ⊗  (in_channels_eff * k_size))
        """

        L = self.fishleg_aux["L"]
        R = self.fishleg_aux["R"]
        A = self.fishleg_aux["A"]

        if self._bias:
            w, b = v
            sw = w.shape
            sb = b.shape
            w = torch.reshape(w, (self.out_channels, -1))
            b = torch.reshape(b, (self.out_channels, 1))
            u = torch.cat([w, b], dim=-1)
        else:
            (w,) = v
            sw = w.shape
            u = torch.reshape(w, (self.out_channels, -1))

        # at this stage, u is out_channels * (in_channels_eff * k_size (perhaps +1))

        u = torch.linalg.multi_dot((R, (A * u), L.T))
        u = A * torch.linalg.multi_dot((R.T, u, L))

        if self._bias:
            return (torch.reshape(u[:, :-1], sw), torch.reshape(u[:, -1], sb))
        else:
            return (torch.reshape(u, sw),)
//  [python]
def partial_deriv(self, quantity, dd, idx=None):
        assert quantity.ndim == 1
        if idx is None:
            return torch.matmul(self._dmat(dd), quantity)
        return torch.matmul(self._dmat(dd)[idx], quantity)
//  [python]
def partial_deriv(self, quantity, dd, idx=None):
        assert quantity.ndim == 1
        if idx is None:
            return torch.linalg.multi_dot((self._dmat(dd), quantity))
        return torch.linalg.multi_dot((self._dmat(dd)[idx], quantity))
//  [python]
def high_order_partial_deriv(self, order, quantity, dd, idx=None):
        if idx is None:
            return torch.mm(*([self._dmat(dd)]*order+[quantity]))
        return torch.mm(
            torch.mm(*([self._dmat(dd)]*order))[idx],
             quantity)
//  [python]
def high_order_partial_deriv(self, order, quantity, dd, idx=None):
        if idx is None:
            return torch.linalg.multi_dot([self._dmat(dd)]*order+[quantity])
        return torch.linalg.multi_dot(
            (torch.linalg.multi_dot([self._dmat(dd)]*order)[idx],
             quantity))
//  [python]
def _apply_periodic_boundary_conditions(
            self, bndry_conds, residual, jac, sol):
        for ii in range(len(bndry_conds)//2):
            if (self._basis_types[ii] == "C" and bndry_conds[2*ii][1] == "P"):
                idx1 = self._bndry_indices[2*ii]
                idx2 = self._bndry_indices[2*ii+1]
                jac[idx1, :] = 0
                jac[idx1, idx1] = 1
                jac[idx1, idx2] = -1
                jac[idx2] = self._dmat(ii//2)[idx1]-self._dmat(ii//2)[idx2]
                residual[idx1] = sol[idx1]-sol[idx2]
                residual[idx2] = (
                    torch.mm(self._dmat(ii//2)[idx1].unsqueeze(0), sol) -
                    torch.mm(self._dmat(ii//2)[idx2].unsqueeze(0), sol))
        return residual, jac
//  [python]
def _apply_periodic_boundary_conditions(
            self, bndry_conds, residual, jac, sol):
        for ii in range(len(bndry_conds)//2):
            if (self._basis_types[ii] == "C" and bndry_conds[2*ii][1] == "P"):
                idx1 = self._bndry_indices[2*ii]
                idx2 = self._bndry_indices[2*ii+1]
                jac[idx1, :] = 0
                jac[idx1, idx1] = 1
                jac[idx1, idx2] = -1
                jac[idx2] = self._dmat(ii//2)[idx1]-self._dmat(ii//2)[idx2]
                residual[idx1] = sol[idx1]-sol[idx2]
                residual[idx2] = (
                    torch.linalg.multi_dot((self._dmat(ii//2)[idx1], sol)) -
                    torch.linalg.multi_dot((self._dmat(ii//2)[idx2], sol)))
        return residual, jac
//  [python]
def partial_deriv(self, quantity, dd, idx=None):
        assert quantity.ndim == 1
        if idx is None:
            return torch.matmul(self._dmat(quantity, dd), quantity)
        return torch.matmul(self._dmat(quantity, dd)[idx], quantity)
//  [python]
def partial_deriv(self, quantity, dd, idx=None):
        assert quantity.ndim == 1
        if idx is None:
            return torch.linalg.multi_dot((self._dmat(quantity, dd), quantity))
        return torch.linalg.multi_dot(
            (self._dmat(quantity, dd)[idx], quantity))
//  [python]
def _apply_neumann_and_robin_boundary_conditions_to_residual(
            self, bndry_conds, residual, sol, flux_jac):
        for ii, bndry_cond in enumerate(bndry_conds):
            if bndry_cond[1] == "N" or bndry_cond[1] == "R":
                if self._basis_types[ii//2] == "F":
                    msg = "Cannot enforce non-periodic boundary conditions "
                    msg += "when using a Fourier basis"
                    raise ValueError(msg)
                idx = self._bndry_indices[ii]
                bndry_vals = bndry_cond[0](self.mesh_pts[:, idx])[:, 0]
                normal_vals = self._bndrys[ii].normals(self.mesh_pts[:, idx])
                flux_jac_vals = flux_jac(idx)
                flux_vals = torch.hstack([
                    torch.matmul(flux_jac_vals[dd], sol)[:, None]
                    for dd in range(len(flux_jac_vals))])
                residual[idx] = self.dot(flux_vals, normal_vals)-bndry_vals
                if bndry_cond[1] == "R":
                    residual[idx] += bndry_cond[2]*sol[idx]
        return residual
//  [python]
def _apply_neumann_and_robin_boundary_conditions_to_residual(
            self, bndry_conds, residual, sol, flux_jac):
        for ii, bndry_cond in enumerate(bndry_conds):
            if bndry_cond[1] == "N" or bndry_cond[1] == "R":
                if self._basis_types[ii//2] == "F":
                    msg = "Cannot enforce non-periodic boundary conditions "
                    msg += "when using a Fourier basis"
                    raise ValueError(msg)
                idx = self._bndry_indices[ii]
                bndry_vals = bndry_cond[0](self.mesh_pts[:, idx])[:, 0]
                normal_vals = self._bndrys[ii].normals(self.mesh_pts[:, idx])
                flux_jac_vals = flux_jac(idx)
                flux_vals = torch.hstack([
                    torch.linalg.multi_dot((flux_jac_vals[dd], sol))[:, None]
                    for dd in range(len(flux_jac_vals))])
                residual[idx] = self.dot(flux_vals, normal_vals)-bndry_vals
                if bndry_cond[1] == "R":
                    residual[idx] += bndry_cond[2]*sol[idx]
        return residual
//  [python]
def _raw_jacobian(self, sol):
        split_sols = self.mesh.split_quantities(sol)
        vel_sols = torch.hstack([s[:, None] for s in split_sols[:-1]])
        vel_forc_vals = self._vel_forc_fun(self.mesh._meshes[0].mesh_pts)
        jac = [[0 for jj in range(self.mesh.nphys_vars+1)]
               for ii in range(len(split_sols))]
        # assumes x and y velocity meshes are the same
        vel_dmats = self.mesh._meshes[0]._dmats
        #[self.mesh._meshes[0]._dmat(dd)
        #for dd in range(self.mesh.nphys_vars)]
        for dd in range(self.mesh.nphys_vars):
            for ii in range(self.mesh.nphys_vars):
                dmat = vel_dmats[ii] # self.mesh._meshes[dd]._dmat(ii)
                jac[dd][dd] += torch.matmul(dmat, dmat)
                if dd != ii:
                    jac[dd][ii] = torch.zeros_like(dmat)
            jac[dd][-1] = -self.mesh._meshes[-1]._dmat(split_sols[-1], dd)
            if self._navier_stokes:
                # 1D
                # residual[0] += v[0]*(D[0]v[0])
                # d residual[0]/dv[0] += diag(D[0]v[dd]) + diag(v[0])*D[0]
                # 2D
                # residual[0] = v[0]*(D[0]v[0]) + v[1]*(D[1]v[0])
                # d residual[0] /dv[0] = diag(D[0]v[0]) + diag(v[0])D[0] + diag(v[1])D[1]
                # d residual[0] /dv[1] = diag(D[1]v[0])
                # residual[1] = v[0]*(D[0]v[1]) + v[1]*(D[1]v[1])
                # d residual[1] /dv[0] = diag(D[0]v[1])
                # d residual[1] /dv[1] = diag(v[0])D[0] + diag(D[1]v[1]) + diag(v[1])D[1]
                for ii in range(self.mesh.nphys_vars):
                    if ii != dd:
                        jac[dd][ii] -= torch.diag(
                            torch.matmul(vel_dmats[ii], split_sols[dd]))
                    jac[dd][dd] -= split_sols[ii][:, None]*vel_dmats[ii]
                jac[dd][dd] -= torch.diag(
                     torch.matmul(vel_dmats[dd], split_sols[dd]))
            jac[dd] = torch.hstack(jac[dd])
        for dd in range(self.mesh.nphys_vars):
            jac[-1][dd] = -self.mesh._meshes[-1]._dmat(split_sols[dd], dd)
        jac[-1][-1] = torch.zeros(
            (self.mesh._meshes[-1].nunknowns, self.mesh._meshes[-1].nunknowns))
        jac[-1] = torch.hstack(jac[-1])
        jac[-1][self._unique_pres_data[0], :] = 0
        jac[-1][self._unique_pres_data[0],
                self.mesh._meshes[0].nunknowns*self.mesh.nphys_vars +
                self._unique_pres_data[0]] = 1
        # Todo reverse sign of residual for time integration
        return torch.vstack(jac)
//  [python]
def _raw_jacobian(self, sol):
        split_sols = self.mesh.split_quantities(sol)
        vel_sols = torch.hstack([s[:, None] for s in split_sols[:-1]])
        vel_forc_vals = self._vel_forc_fun(self.mesh._meshes[0].mesh_pts)
        jac = [[0 for jj in range(self.mesh.nphys_vars+1)]
               for ii in range(len(split_sols))]
        # assumes x and y velocity meshes are the same
        vel_dmats = self.mesh._meshes[0]._dmats
        #[self.mesh._meshes[0]._dmat(dd)
        #for dd in range(self.mesh.nphys_vars)]
        for dd in range(self.mesh.nphys_vars):
            for ii in range(self.mesh.nphys_vars):
                dmat = vel_dmats[ii] # self.mesh._meshes[dd]._dmat(ii)
                jac[dd][dd] += multi_dot((dmat, dmat))
                if dd != ii:
                    jac[dd][ii] = torch.zeros_like(dmat)
            jac[dd][-1] = -self.mesh._meshes[-1]._dmat(split_sols[-1], dd)
            if self._navier_stokes:
                # 1D
                # residual[0] += v[0]*(D[0]v[0])
                # d residual[0]/dv[0] += diag(D[0]v[dd]) + diag(v[0])*D[0]
                # 2D
                # residual[0] = v[0]*(D[0]v[0]) + v[1]*(D[1]v[0])
                # d residual[0] /dv[0] = diag(D[0]v[0]) + diag(v[0])D[0] + diag(v[1])D[1]
                # d residual[0] /dv[1] = diag(D[1]v[0])
                # residual[1] = v[0]*(D[0]v[1]) + v[1]*(D[1]v[1])
                # d residual[1] /dv[0] = diag(D[0]v[1])
                # d residual[1] /dv[1] = diag(v[0])D[0] + diag(D[1]v[1]) + diag(v[1])D[1]
                for ii in range(self.mesh.nphys_vars):
                    if ii != dd:
                        jac[dd][ii] -= torch.diag(
                            multi_dot((vel_dmats[ii], split_sols[dd])))
                    jac[dd][dd] -= split_sols[ii][:, None]*vel_dmats[ii]
                jac[dd][dd] -= torch.diag(
                     multi_dot((vel_dmats[dd], split_sols[dd])))
            jac[dd] = torch.hstack(jac[dd])
        for dd in range(self.mesh.nphys_vars):
            jac[-1][dd] = -self.mesh._meshes[-1]._dmat(split_sols[dd], dd)
        jac[-1][-1] = torch.zeros(
            (self.mesh._meshes[-1].nunknowns, self.mesh._meshes[-1].nunknowns))
        jac[-1] = torch.hstack(jac[-1])
        jac[-1][self._unique_pres_data[0], :] = 0
        jac[-1][self._unique_pres_data[0],
                self.mesh._meshes[0].nunknowns*self.mesh.nphys_vars +
                self._unique_pres_data[0]] = 1
        # Todo reverse sign of residual for time integration
        return torch.vstack(jac)
//  [python]
def _apply_neumann_and_robin_boundary_conditions(
            self, bndry_conds, residual, jac, sol, flux_jac):
        for ii, bndry_cond in enumerate(bndry_conds):
            if bndry_cond[1] != "N" and bndry_cond[1] != "R":
                continue
            idx = self._bndry_indices[ii]
            mesh_pts_idx = self._bndry_slice(self.mesh_pts, idx, 1)
            if self._normal_vals[ii] is None:
                self._normal_vals[ii] = self._bndrys[ii].normals(mesh_pts_idx)
            if not self._flux_islinear or self._flux_normal_vals[ii] is None:
                flux_jac_vals = flux_jac(idx)
                flux_normal_vals = [
                    self._normal_vals[ii][:, dd:dd+1]*flux_jac_vals[dd]
                    for dd in range(self.nphys_vars)]
                if self._flux_islinear:
                    self._flux_normal_vals[ii] = flux_normal_vals
            else:
                flux_normal_vals = self._flux_normal_vals[ii]
            # (D2*u)*n2+D2*u*n2
            jac[idx] = sum(flux_normal_vals)
            bndry_vals = bndry_cond[0](mesh_pts_idx)[:, 0]

            # residual[idx] = torch.linalg.multi_dot((jac[idx], sol))-bndry_vals
            residual[idx] = torch.mm(self._bndry_slice(jac, idx, 0), sol)-bndry_vals
            if bndry_cond[1] == "R":
                jac[idx, idx] += bndry_cond[2]
                # residual[idx] += bndry_cond[2]*sol[idx]
                residual[idx] += bndry_cond[2]*self._bndry_slice(sol, idx, 0)

        return residual, jac
//  [python]
def _apply_neumann_and_robin_boundary_conditions(
            self, bndry_conds, residual, jac, sol, flux_jac):
        for ii, bndry_cond in enumerate(bndry_conds):
            if bndry_cond[1] != "N" and bndry_cond[1] != "R":
                continue
            idx = self._bndry_indices[ii]
            mesh_pts_idx = self._bndry_slice(self.mesh_pts, idx, 1)
            if self._normal_vals[ii] is None:
                self._normal_vals[ii] = self._bndrys[ii].normals(mesh_pts_idx)
            if not self._flux_islinear or self._flux_normal_vals[ii] is None:
                flux_jac_vals = flux_jac(idx)
                flux_normal_vals = [
                    self._normal_vals[ii][:, dd:dd+1]*flux_jac_vals[dd]
                    for dd in range(self.nphys_vars)]
                if self._flux_islinear:
                    self._flux_normal_vals[ii] = flux_normal_vals
            else:
                flux_normal_vals = self._flux_normal_vals[ii]
            # (D2*u)*n2+D2*u*n2
            jac[idx] = sum(flux_normal_vals)
            bndry_vals = bndry_cond[0](mesh_pts_idx)[:, 0]

            # residual[idx] = torch.linalg.multi_dot((jac[idx], sol))-bndry_vals
            residual[idx] = torch.linalg.multi_dot(
                (self._bndry_slice(jac, idx, 0), sol))-bndry_vals
            if bndry_cond[1] == "R":
                jac[idx, idx] += bndry_cond[2]
                # residual[idx] += bndry_cond[2]*sol[idx]
                residual[idx] += bndry_cond[2]*self._bndry_slice(sol, idx, 0)

        return residual, jac
//  [python]
def simulate(self, obj, snr_db=20):
        """
        Simulate the mask measurement of an image. Apply left and right convolution matrices,
        add noise and return the measurement.

        Parameters
        ----------
        obj: :py:class:`~numpy.ndarray`
            Image to simulate.
        snr_db: float, optional
            Signal-to-noise ratio (dB) of the simulated measurement. Default is 20 dB.
        """
        assert len(obj.shape) == 3, "Object should be a 3D array (HxWxC) even if grayscale."

        # Get convolution matrices
        P, Q = self.get_conv_matrices(obj.shape)

        # Convolve image
        n_channels = obj.shape[-1]

        if torch_available and isinstance(obj, torch.Tensor):
            P = torch.from_numpy(P).float()
            Q = torch.from_numpy(Q).float()
            meas = torch.dstack(
                [torch.mm(torch.mm(P, obj[:, :, c]), Q.T) for c in range(n_channels)]
            ).float()
        else:
            meas = np.dstack([multi_dot([P, obj[:, :, c], Q.T]) for c in range(n_channels)])

        # Add noise
        if snr_db is not None:
            meas = add_shot_noise(meas, snr_db=snr_db)

        if torch_available and isinstance(obj, torch.Tensor):
            meas = meas.to(obj)

        return meas
//  [python]
def simulate(self, obj, snr_db=20):
        """
        Simulate the mask measurement of an image. Apply left and right convolution matrices,
        add noise and return the measurement.

        Parameters
        ----------
        obj: :py:class:`~numpy.ndarray`
            Image to simulate.
        snr_db: float, optional
            Signal-to-noise ratio (dB) of the simulated measurement. Default is 20 dB.
        """
        assert len(obj.shape) == 3, "Object should be a 3D array (HxWxC) even if grayscale."

        # Get convolution matrices
        P, Q = self.get_conv_matrices(obj.shape)

        # Convolve image
        n_channels = obj.shape[-1]

        if torch_available and isinstance(obj, torch.Tensor):
            P = torch.from_numpy(P).float()
            Q = torch.from_numpy(Q).float()
            meas = torch.dstack(
                [torch.linalg.multi_dot([P, obj[:, :, c], Q.T]) for c in range(n_channels)]
            ).float()
        else:
            meas = np.dstack([multi_dot([P, obj[:, :, c], Q.T]) for c in range(n_channels)])

        # Add noise
        if snr_db is not None:
            meas = add_shot_noise(meas, snr_db=snr_db)

        if torch_available and isinstance(obj, torch.Tensor):
            meas = meas.to(obj)

        return meas
//  [python]
def test_multi_dot_errors(self, device, dtype):
        def check(tensors, out, msg):
            with self.assertRaisesRegex(RuntimeError, msg):
                t.multi_dot(tensors, out=out)

        a = make_tensor(2, device, dtype)

        check([], None, "expected at least 2 tensors")
        check([a], None, "expected at least 2 tensors")

        check([torch.tensor(1, device=device, dtype=dtype), a], None, "the first tensor must be 1D or 2D")
        check([a, torch.tensor(1, device=device, dtype=dtype)], None, "the last tensor must be 1D or 2D")

        check([a, a, a], None, "tensor 1 must be 2D")
        check([a, make_tensor((2, 2, 2), device, dtype), a], None, "tensor 1 must be 2D")

        check([a, make_tensor(2, device, torch.double)], None, "all tensors must have be the same dtype")
        check([a, a], torch.empty(0, device=device, dtype=torch.double), "expected out tensor to have dtype")

        if self.device_type == 'cuda':
            check([a, make_tensor(2, 'cpu', dtype)], None, "all tensors must be on the same device")
            check([a, a], torch.empty(0, dtype=dtype), "expected out tensor to be on device")

        check([a, make_tensor(3, device, dtype)], None, "cannot be multiplied")
        check([a, make_tensor((3, 2), device, dtype), a], None, "cannot be multiplied")
//  [python]
def test_multi_dot_errors(self, device, dtype):
        def check(tensors, out, msg):
            with self.assertRaisesRegex(RuntimeError, msg):
                torch.linalg.multi_dot(tensors, out=out)

        a = make_tensor(2, device, dtype)

        check([], None, "expected at least 2 tensors")
        check([a], None, "expected at least 2 tensors")

        check([torch.tensor(1, device=device, dtype=dtype), a], None, "the first tensor must be 1D or 2D")
        check([a, torch.tensor(1, device=device, dtype=dtype)], None, "the last tensor must be 1D or 2D")

        check([a, a, a], None, "tensor 1 must be 2D")
        check([a, make_tensor((2, 2, 2), device, dtype), a], None, "tensor 1 must be 2D")

        check([a, make_tensor(2, device, torch.double)], None, "all tensors must have be the same dtype")
        check([a, a], torch.empty(0, device=device, dtype=torch.double), "expected out tensor to have dtype")

        if self.device_type == 'cuda':
            check([a, make_tensor(2, 'cpu', dtype)], None, "all tensors must be on the same device")
            check([a, a], torch.empty(0, dtype=dtype), "expected out tensor to be on device")

        check([a, make_tensor(3, device, dtype)], None, "cannot be multiplied")
        check([a, make_tensor((3, 2), device, dtype), a], None, "cannot be multiplied")
//  [python]
def compute_gradient(self, set_param_values, param_vals, return_obj=False,
                         **newton_kwargs):
        # use detach so that fwd_sol is not part of AD-graph
        set_param_values(self._fwd_solver.physics, param_vals.detach())
        fwd_sol = self._fwd_solver.solve(**newton_kwargs)
        adj_sol = self.solve_adjoint(fwd_sol, param_vals.detach())

        param_vals_copy = torch.clone(param_vals)
        fwd_sol_copy = torch.clone(fwd_sol)
        if self._dqdp is None:
            functional_vals = self._functional(
                fwd_sol_copy.requires_grad_(True),
                param_vals_copy).requires_grad_(True)
            functional_vals.backward()
            dqdp = param_vals_copy.grad
            if dqdp is None:
                dqdp = 0
        else:
            functional_vals = self._functional(fwd_sol_copy, param_vals_copy)
            dqdp = self._dqdp(fwd_sol_copy, param_vals_copy)

        if self._dRdp is None:
            dRdp = torch.autograd.functional.jacobian(
                partial(self._parameterized_residual, fwd_sol,
                        set_param_values), param_vals, strict=True)
        else:
            dRdp = self._dRdp(
                self._fwd_solver.physics, fwd_sol_copy, param_vals)
        # print("dqdp", dqdp)
        # print(dRdp, 'dRdp')
        # print(adj_sol, 'asol')
        # print(fwd_sol, 'fsol')
        grad = dqdp+torch.matmul(adj_sol[:,None], dRdp)
        if not return_obj:
            return grad
        return functional_vals.detach(), grad.detach()
//  [python]
def compute_gradient(self, set_param_values, param_vals, return_obj=False,
                         **newton_kwargs):
        # use detach so that fwd_sol is not part of AD-graph
        set_param_values(self._fwd_solver.physics, param_vals.detach())
        fwd_sol = self._fwd_solver.solve(**newton_kwargs)
        adj_sol = self.solve_adjoint(fwd_sol, param_vals.detach())

        param_vals_copy = torch.clone(param_vals)
        fwd_sol_copy = torch.clone(fwd_sol)
        if self._dqdp is None:
            functional_vals = self._functional(
                fwd_sol_copy.requires_grad_(True),
                param_vals_copy).requires_grad_(True)
            functional_vals.backward()
            dqdp = param_vals_copy.grad
            if dqdp is None:
                dqdp = 0
        else:
            functional_vals = self._functional(fwd_sol_copy, param_vals_copy)
            dqdp = self._dqdp(fwd_sol_copy, param_vals_copy)

        if self._dRdp is None:
            dRdp = torch.autograd.functional.jacobian(
                partial(self._parameterized_residual, fwd_sol,
                        set_param_values), param_vals, strict=True)
        else:
            dRdp = self._dRdp(
                self._fwd_solver.physics, fwd_sol_copy, param_vals)
        # print("dqdp", dqdp)
        # print(dRdp, 'dRdp')
        # print(adj_sol, 'asol')
        # print(fwd_sol, 'fsol')
        grad = dqdp+torch.linalg.multi_dot((adj_sol[None, :], dRdp))
        if not return_obj:
            return grad
        return functional_vals.detach(), grad.detach()
//  [python]
def warp(
        pos1,
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2
):
    device = pos1.device

    Z1, pos1, ids = interpolate_depth(pos1, depth1)

    # COLMAP convention
    u1 = pos1[1, :] + bbox1[1] + .5
    v1 = pos1[0, :] + bbox1[0] + .5

    X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
    Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

    XYZ1_hom = torch.cat([
        X1.view(1, -1),
        Y1.view(1, -1),
        Z1.view(1, -1),
        torch.ones(1, Z1.size(0), device=device)
    ], dim=0)
    XYZ2_hom = torch.mm(pose2, torch.mm(torch.inverse(pose1), XYZ1_hom))
    XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

    uv2_hom = torch.matmul(intrinsics2, XYZ2)
    uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

    u2 = uv2[0, :] - bbox2[1] - .5
    v2 = uv2[1, :] - bbox2[0] - .5
    uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

    annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

    ids = ids[new_ids]
    pos1 = pos1[:, new_ids]
    estimated_depth = XYZ2[2, new_ids]

    inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

    ids = ids[inlier_mask]
    if ids.size(0) == 0:
        raise EmptyTensorError

    pos2 = pos2[:, inlier_mask]
    pos1 = pos1[:, inlier_mask]

    return pos1, pos2, ids
//  [python]
def warp(
        pos1,
        depth1, intrinsics1, pose1, bbox1,
        depth2, intrinsics2, pose2, bbox2
):
    device = pos1.device

    Z1, pos1, ids = interpolate_depth(pos1, depth1)

    # COLMAP convention
    u1 = pos1[1, :] + bbox1[1] + .5
    v1 = pos1[0, :] + bbox1[0] + .5

    X1 = (u1 - intrinsics1[0, 2]) * (Z1 / intrinsics1[0, 0])
    Y1 = (v1 - intrinsics1[1, 2]) * (Z1 / intrinsics1[1, 1])

    XYZ1_hom = torch.cat([
        X1.view(1, -1),
        Y1.view(1, -1),
        Z1.view(1, -1),
        torch.ones(1, Z1.size(0), device=device)
    ], dim=0)
    # XYZ2_hom = torch.chain_matmul(pose2, torch.inverse(pose1), XYZ1_hom)
    XYZ2_hom = torch.linalg.multi_dot((pose2, torch.inverse(pose1), XYZ1_hom))
    XYZ2 = XYZ2_hom[: -1, :] / XYZ2_hom[-1, :].view(1, -1)

    uv2_hom = torch.matmul(intrinsics2, XYZ2)
    uv2 = uv2_hom[: -1, :] / uv2_hom[-1, :].view(1, -1)

    u2 = uv2[0, :] - bbox2[1] - .5
    v2 = uv2[1, :] - bbox2[0] - .5
    uv2 = torch.cat([u2.view(1, -1),  v2.view(1, -1)], dim=0)

    annotated_depth, pos2, new_ids = interpolate_depth(uv_to_pos(uv2), depth2)

    ids = ids[new_ids]
    pos1 = pos1[:, new_ids]
    estimated_depth = XYZ2[2, new_ids]

    inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05

    ids = ids[inlier_mask]
    if ids.size(0) == 0:
        raise EmptyTensorError

    pos2 = pos2[:, inlier_mask]
    pos1 = pos1[:, inlier_mask]

    return pos1, pos2, ids
//  [python]
def Qv(self, v: Tuple, full=False) -> Tuple:
        Sk = self.fishleg_aux["Sk"]
        Sq = self.fishleg_aux["Sq"]
        Sv = self.fishleg_aux["Sv"]
        So = self.fishleg_aux["So"]

        Uk = Sk * torch.transpose(torch.cat([v[0], v[1][:, None]], dim=-1), -1, -2)
        Uq = Sq * torch.cat([v[2], v[3][:, None]], dim=-1)
        Uv = Sv * torch.cat([v[4], v[5][:, None]], dim=-1)
        Uo = So * torch.cat([v[6], v[7][:, None]], dim=-1)

        zkq = torch.mm(self.fishleg_aux["A"].T, torch.mm(Uk, self.fishleg_aux["Q"])) + torch.mm(self.fishleg_aux["K"], torch.mm(Uq, self.fishleg_aux["B"].T))

        zvo = torch.mm(self.fishleg_aux["O"], torch.mm(Uv, self.fishleg_aux["C"].T)) + torch.mm(self.fishleg_aux["D"].T, torch.mm(Uo, self.fishleg_aux["V"]))

        L = self.fishleg_aux["L"]
        U = self.fishleg_aux["U"]
        R = self.fishleg_aux["R"]

        zkq = torch.mm(torch.mm(R, R.T), torch.mm(zkq, torch.mm(L, L.T)))

        zvo = torch.mm(torch.mm(U, U.T), torch.mm(zvo, torch.mm(L, L.T)))

        Vk = torch.transpose(
            Sk
            * torch.mm(self.fishleg_aux["A"], torch.mm(zkq, self.fishleg_aux["Q"].T)),
            -1,
            -2,
        )

        Vq = Sq * torch.mm(self.fishleg_aux["K"].T, torch.mm(zkq, self.fishleg_aux["B"]))

        Vv = Sv * torch.mm(self.fishleg_aux["O"].T, torch.mm(zvo, self.fishleg_aux["C"]))

        Vo = So * torch.mm(self.fishleg_aux["D"], torch.mm(zvo, self.fishleg_aux["V"].T))

        return (
            Vk[:, :-1],
            Vk[:, -1],
            Vq[:, :-1],
            Vq[:, -1],
            Vv[:, :-1],
            Vv[:, -1],
            Vo[:, :-1],
            Vo[:, -1],
        )
//  [python]
def Qv(self, v: Tuple, full=False) -> Tuple:
        Sk = self.fishleg_aux["Sk"]
        Sq = self.fishleg_aux["Sq"]
        Sv = self.fishleg_aux["Sv"]
        So = self.fishleg_aux["So"]

        Uk = Sk * torch.transpose(torch.cat([v[0], v[1][:, None]], dim=-1), -1, -2)
        Uq = Sq * torch.cat([v[2], v[3][:, None]], dim=-1)
        Uv = Sv * torch.cat([v[4], v[5][:, None]], dim=-1)
        Uo = So * torch.cat([v[6], v[7][:, None]], dim=-1)

        zkq = torch.linalg.multi_dot(
            (self.fishleg_aux["A"].T, Uk, self.fishleg_aux["Q"])
        ) + torch.linalg.multi_dot((self.fishleg_aux["K"], Uq, self.fishleg_aux["B"].T))

        zvo = torch.linalg.multi_dot(
            (self.fishleg_aux["O"], Uv, self.fishleg_aux["C"].T)
        ) + torch.linalg.multi_dot((self.fishleg_aux["D"].T, Uo, self.fishleg_aux["V"]))

        L = self.fishleg_aux["L"]
        U = self.fishleg_aux["U"]
        R = self.fishleg_aux["R"]

        zkq = torch.linalg.multi_dot((R, R.T, zkq, L, L.T))

        zvo = torch.linalg.multi_dot((U, U.T, zvo, L, L.T))

        Vk = torch.transpose(
            Sk
            * torch.linalg.multi_dot(
                (self.fishleg_aux["A"], zkq, self.fishleg_aux["Q"].T)
            ),
            -1,
            -2,
        )

        Vq = Sq * torch.linalg.multi_dot(
            (self.fishleg_aux["K"].T, zkq, self.fishleg_aux["B"])
        )

        Vv = Sv * torch.linalg.multi_dot(
            (self.fishleg_aux["O"].T, zvo, self.fishleg_aux["C"])
        )

        Vo = So * torch.linalg.multi_dot(
            (self.fishleg_aux["D"], zvo, self.fishleg_aux["V"].T)
        )

        return (
            Vk[:, :-1],
            Vk[:, -1],
            Vq[:, :-1],
            Vq[:, -1],
            Vv[:, :-1],
            Vv[:, -1],
            Vo[:, :-1],
            Vo[:, -1],
        )
//  [python]
def compute_gradient(self, init_sol, init_time, final_time,
                         set_param_values, param_vals, **newton_kwargs):
        if init_sol.ndim == 2:
            assert init_sol.shape[1] == 1
            init_sol = init_sol[:, 0]
        fwd_sols, times, qoi, dqdp = self.compute_qoi(
            init_sol, init_time, final_time,
            set_param_values, param_vals, **newton_kwargs)

        adj_sols = self.solve_adjoint(fwd_sols, param_vals, times)

        phi_dRdp = 0
        for ii in range(fwd_sols.shape[1]):
            if ii > 0:
                deltat = times[ii]-times[ii-1]
            else:
                deltat = None
            dRdp_ii = torch.autograd.functional.jacobian(
                partial(self._parameterized_transient_residual,
                        fwd_sols[:, ii], times[ii], set_param_values,
                        init_sol, init_time, deltat), param_vals)#, strict=True)
            # print(torch.autograd.functional.jacobian(
            #     partial(self._parameterized_transient_residual,
            #             time=times[ii], set_param_values=set_param_values,
            #             init_sol=init_sol, init_time=init_time,
            #             deltat=deltat, param_vals=param_vals), fwd_sols[:, ii]))
            # print(ii, dRdp_ii, 'dRdp')
            # print(adj_sols[:, ii:ii+1].T)
            phi_dRdp += torch.matmul(adj_sols[:, ii:ii+1].T, dRdp_ii)
        grad = dqdp + phi_dRdp
        return np.atleast_1d(qoi), grad.numpy()
//  [python]
def compute_gradient(self, init_sol, init_time, final_time,
                         set_param_values, param_vals, **newton_kwargs):
        if init_sol.ndim == 2:
            assert init_sol.shape[1] == 1
            init_sol = init_sol[:, 0]
        fwd_sols, times, qoi, dqdp = self.compute_qoi(
            init_sol, init_time, final_time,
            set_param_values, param_vals, **newton_kwargs)

        adj_sols = self.solve_adjoint(fwd_sols, param_vals, times)

        phi_dRdp = 0
        for ii in range(fwd_sols.shape[1]):
            if ii > 0:
                deltat = times[ii]-times[ii-1]
            else:
                deltat = None
            dRdp_ii = torch.autograd.functional.jacobian(
                partial(self._parameterized_transient_residual,
                        fwd_sols[:, ii], times[ii], set_param_values,
                        init_sol, init_time, deltat), param_vals)#, strict=True)
            # print(torch.autograd.functional.jacobian(
            #     partial(self._parameterized_transient_residual,
            #             time=times[ii], set_param_values=set_param_values,
            #             init_sol=init_sol, init_time=init_time,
            #             deltat=deltat, param_vals=param_vals), fwd_sols[:, ii]))
            # print(ii, dRdp_ii, 'dRdp')
            # print(adj_sols[:, ii:ii+1].T)
            phi_dRdp += torch.linalg.multi_dot(
                (adj_sols[:, ii:ii+1].T, dRdp_ii))
        grad = dqdp + phi_dRdp
        return np.atleast_1d(qoi), grad.numpy()
//  [python]
def get_y_neg_log_likelihood(self, Y, X):
        """
        L_y = D/2*log(|Σ_y|) + 1/2*trace(vec(Y)*vec(Y)^T*Σ_y^-1)
        """
        K_y = self.observationGP_kernel(X, X)

        # get Ky_inv
        U, info = torch.linalg.cholesky_ex(K_y, upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        # get K_D_inv
        U, info = torch.linalg.cholesky_ex(self.K_D, upper=True)
        U_inv = torch.inverse(U)
        K_D_inv = U_inv @ U_inv.t()
        # get Sigma_y_inv
        Sigma_y_inv = torch.kron(Ky_inv, K_D_inv)
        # get log det of Sigma_y
        log_det_K_y = torch.logdet(K_y)
        log_det_K_D = torch.logdet(self.K_D)
        log_det_Sigma_y = K_y.size()[0] * log_det_K_D + self.K_D.size()[0] * log_det_K_y
        # TODO: add noise term?

        return self.D / 2 * log_det_Sigma_y + 1 / 2 * torch.trace(torch.mm(torch.mm(Y.view(-1,1), Y.view(-1,1).t()), Sigma_y_inv))
//  [python]
def get_y_neg_log_likelihood(self, Y, X):
        """
        L_y = D/2*log(|Σ_y|) + 1/2*trace(vec(Y)*vec(Y)^T*Σ_y^-1)
        """
        K_y = self.observationGP_kernel(X, X)

        # get Ky_inv
        U, info = torch.linalg.cholesky_ex(K_y, upper=True)
        U_inv = torch.inverse(U)
        Ky_inv = U_inv @ U_inv.t()
        # get K_D_inv
        U, info = torch.linalg.cholesky_ex(self.K_D, upper=True)
        U_inv = torch.inverse(U)
        K_D_inv = U_inv @ U_inv.t()
        # get Sigma_y_inv
        Sigma_y_inv = torch.kron(Ky_inv, K_D_inv)
        # get log det of Sigma_y
        log_det_K_y = torch.logdet(K_y)
        log_det_K_D = torch.logdet(self.K_D)
        log_det_Sigma_y = K_y.size()[0] * log_det_K_D + self.K_D.size()[0] * log_det_K_y
        # TODO: add noise term?

        return self.D / 2 * log_det_Sigma_y + 1 / 2 * torch.trace(torch.linalg.multi_dot([Y.view(-1,1), Y.view(-1,1).t(), Sigma_y_inv]))
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes, noncontiguous=False):
            tensors = [make_tensor(shape, device, dtype, noncontiguous=noncontiguous) 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], t.mm(tensors[2], tensors[3]))) if len(tensors) == 4 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])

        # test noncontiguous input
        check([3, 2], [2, 2], [2, 3], [3, 4], noncontiguous=True)
        check([15, 5], [5, 10], [10, 20], [20, 25], noncontiguous=True)
//  [python]
def test_multi_dot(self, device, dtype):
        def check(*shapes, noncontiguous=False):
            tensors = [make_tensor(shape, device, dtype, noncontiguous=noncontiguous) 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])

        # test noncontiguous input
        check([3, 2], [2, 2], [2, 3], [3, 4], noncontiguous=True)
        check([15, 5], [5, 10], [10, 20], [20, 25], noncontiguous=True)
//  [python]
def multi_dot(
    x: Sequence[torch.Tensor],
    /,
    *,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    res = x[0]
    for i in range(1, len(x)):
        res = torch.matmul(res, x[i])
    return res
//  [python]
def multi_dot(
    x: Sequence[torch.Tensor],
    /,
    *,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    return torch.linalg.multi_dot(x, out=out)
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin):
        """
        L_x = Q/2*log(|Σ_x|) + 1/2*trace(vec(Xout)*vec(Xout)^T*Σ_x^-1)
        """
        K_x = self.dynamicGP_kernel(Xin, Xin)
        # get Kx_inv
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        # get K_Q_inv
        U, info = torch.linalg.cholesky_ex(self.K_Q, upper=True)
        U_inv = torch.inverse(U)
        K_Q_inv = U_inv @ U_inv.t()
        # get Sigma_x_inv
        Sigma_x_inv = torch.kron(Kx_inv, K_Q_inv)
        # get log det of Sigma_y
        log_det_K_x = torch.logdet(K_x)
        log_det_K_Q = torch.logdet(self.K_Q)
        log_det_Sigma_x = K_x.size()[0] * log_det_K_Q + self.K_Q.size()[0] * log_det_K_x
        # TODO: add noise term?

        return self.Q / 2 * log_det_Sigma_x + 1 / 2 * torch.trace(torch.mm(torch.mm(Xout.view(-1, 1), Xout.view(-1, 1).t()), Sigma_x_inv))
//  [python]
def get_x_neg_log_likelihood(self, Xout, Xin):
        """
        L_x = Q/2*log(|Σ_x|) + 1/2*trace(vec(Xout)*vec(Xout)^T*Σ_x^-1)
        """
        K_x = self.dynamicGP_kernel(Xin, Xin)
        # get Kx_inv
        U, info = torch.linalg.cholesky_ex(K_x, upper=True)
        U_inv = torch.inverse(U)
        Kx_inv = U_inv @ U_inv.t()
        # get K_Q_inv
        U, info = torch.linalg.cholesky_ex(self.K_Q, upper=True)
        U_inv = torch.inverse(U)
        K_Q_inv = U_inv @ U_inv.t()
        # get Sigma_x_inv
        Sigma_x_inv = torch.kron(Kx_inv, K_Q_inv)
        # get log det of Sigma_y
        log_det_K_x = torch.logdet(K_x)
        log_det_K_Q = torch.logdet(self.K_Q)
        log_det_Sigma_x = K_x.size()[0] * log_det_K_Q + self.K_Q.size()[0] * log_det_K_x
        # TODO: add noise term?

        return self.Q / 2 * log_det_Sigma_x + 1 / 2 * torch.trace(torch.linalg.multi_dot([Xout.view(-1, 1), Xout.view(-1, 1).t(), Sigma_x_inv]))
//  [python]
def build_graph(self, config):
        inputs = []
        for i in range(len(config.inputs_shape)):
            input_i = self.variable(
                name='input_' + str(i),
                shape=config.inputs_shape[i],
                dtype=config.inputs_dtype[i])
            inputs.append(input_i)
        result = torch.mm(*inputs)

        self.feed_list = inputs
        self.fetch_list = [result]
//  [python]
def build_graph(self, config):
        inputs = []
        for i in range(len(config.inputs_shape)):
            input_i = self.variable(
                name='input_' + str(i),
                shape=config.inputs_shape[i],
                dtype=config.inputs_dtype[i])
            inputs.append(input_i)
        result = torch.linalg.multi_dot(inputs)

        self.feed_list = inputs
        self.fetch_list = [result]
//  [python]
def test_slogdet_sign(self):
        a = torch.randn(3, 3, dtype=torch.double, requires_grad=True)
        s, logdet = a.slogdet()

        # test that sign should not require grad
        self.assertFalse(s.requires_grad)

        # test that backward through computation involving sign works
        def sign_mul_logdet(mat):
            s, logdet = mat.slogdet()
            return s * logdet

        u, s, v = a.detach().svd()
        s.abs_().clamp_(0.0001)
        for sign in (-1, 1):
            s[-1] = sign
            mat = torch.mm(torch.mm(u, s.diag()), v.t()).requires_grad_()
            gradcheck(sign_mul_logdet, mat)
            gradgradcheck(sign_mul_logdet, mat)
//  [python]
def test_slogdet_sign(self):
        a = torch.randn(3, 3, dtype=torch.double, requires_grad=True)
        s, logdet = a.slogdet()

        # test that sign should not require grad
        self.assertFalse(s.requires_grad)

        # test that backward through computation involving sign works
        def sign_mul_logdet(mat):
            s, logdet = mat.slogdet()
            return s * logdet

        u, s, v = a.detach().svd()
        s.abs_().clamp_(0.0001)
        for sign in (-1, 1):
            s[-1] = sign
            mat = torch.linalg.multi_dot([u, s.diag(), v.t()]).requires_grad_()
            gradcheck(sign_mul_logdet, mat)
            gradgradcheck(sign_mul_logdet, mat)
//  [python]
def forward_observationGP(self, Xstar, flg_noise=False):
        """
        inference of observationGP, map X to Y
        :param Xstar:       X
        :param flg_noise:   whether add noise
        :return:            Y
        """
        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):
        """
        inference of observationGP, map X to Y
        :param Xstar:       X
        :param flg_noise:   whether add noise
        :return:            Y
        """
        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 test_vmap_multi_dot_failure_1D_input(self):
        # special exception for first and last tensors so making giving 3 items avoids special cases
        inputs = (torch.randn(3, 3), torch.randn(3), torch.randn(3, 3))
        with self.assertRaisesRegex(RuntimeError, "tensor 1 must be 2D but got 1D"):
            torch.matmul(inputs[0], torch.matmul(inputs[1], inputs[2]))

        # square inputs are more likely to pass linalg checks
        inputs = tuple(i.expand(i.shape[0], i.shape[0]) for i in inputs)
        with self.assertRaisesRegex(RuntimeError, "tensor 1 must be 2D but got 1D"):
            return vmap(lambda x: torch.matmul(x[0], torch.matmul(x[1], x[2])))(inputs)
//  [python]
def test_vmap_multi_dot_failure_1D_input(self):
        # special exception for first and last tensors so making giving 3 items avoids special cases
        inputs = (torch.randn(3, 3), torch.randn(3), torch.randn(3, 3))
        with self.assertRaisesRegex(RuntimeError, "tensor 1 must be 2D but got 1D"):
            torch.linalg.multi_dot(inputs)

        # square inputs are more likely to pass linalg checks
        inputs = tuple(i.expand(i.shape[0], i.shape[0]) for i in inputs)
        with self.assertRaisesRegex(RuntimeError, "tensor 1 must be 2D but got 1D"):
            return vmap(torch.linalg.multi_dot)(inputs)
//  [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 = (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