code: def _optimize() -> str:
      """Return contraction order for `np.einsum` based on platform.
  
      Introduced after https://github.com/google/jax/pull/7512 since TPU seems to
      be more precise in `greeedy` mode.
      """
      return "greedy" if jax.default_backend() == "tpu" else "optimal"
  
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 6 for calling a global function with example usage: _optimize() and returns: unknown