code: def _is_on_cpu(x: PyTree) -> bool:
      def _arr_is_on_cpu(x: jnp.ndarray) -> bool:
          # TODO(romann): revisit when https://github.com/google/jax/issues/1431 and
          # https://github.com/google/jax/issues/1432 are fixed.
          if hasattr(x, "addressable_shards"):
              # device_buffer is deprecated, so try addressable_shards first.
              return (
                  "cpu" in str(x.addressable_shards[0].device).lower()
              )
          elif hasattr(x, "device_buffer"):
              return "cpu" in str(x.device_buffer.device()).lower()
  
          if isinstance(x, (np.ndarray, jnp.ndarray)):
              return True
  
          raise NotImplementedError(type(x))
  
      return tree_all(tree_map(_arr_is_on_cpu, x))
  
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 18 for calling a global function with example usage: _is_on_cpu(x) and returns: tree_all, unknown, True, unknown