code: def _check_inputs(
      fx_train_or_state_0: Union[ArrayOrScalar, ODEState],
      fx_test_0: ArrayOrScalar,
      k_test_train: Optional[jnp.ndarray],
  ):
      if isinstance(fx_train_or_state_0, ODEState):
          if fx_test_0 is not None:
              raise ValueError(
                  "`fx_test_0` is included in `ODEState` and must be set "
                  "to `None`."
              )
  
          fx_train_0 = fx_train_or_state_0.fx_train
          fx_test_0 = fx_train_or_state_0.fx_test
  
      else:
          fx_train_0 = fx_train_or_state_0
  
      if fx_train_0 is None and fx_test_0 is None:
          raise ValueError(
              "Both `fx_train_0` and `fx_test_0` are `None`, i.e. no "
              "predictions will be computed."
          )
      if fx_test_0 is not None and k_test_train is None:
          raise ValueError(
              "To get predictions on the test set, please provide "
              "`k_test_train` kernel to the parent function."
          )
  
imports: import collections
  from functools import lru_cache
  from typing import Any, Callable, Generator, Iterable, NamedTuple, Optional, Protocol, Union
  import jax
  from jax import grad
  from jax.experimental import ode
  import jax.numpy
  import jax.scipy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  import scipy
  from .utils import dataclasses
  from .utils import utils
  from .utils.typing import Axes
  from .utils.typing import Get
  from .utils.typing import KernelFn
label_desc: Use image node_id 14 for calling a global function with example usage: _check_inputs(fx_train_or_state_0, fx_test_0, k_test_train) without return types