code: def get_runnable(
      bento_model: bentoml.Model,
  ) -> t.Type[bentoml.Runnable]:
      """Private API: use :obj:`~bentoml.Model.to_runnable` instead."""
      partial_kwargs: dict[
          str, t.Any
      ] = bento_model.info.options.partial_kwargs
  
      class FlaxRunnable(bentoml.Runnable):
          SUPPORTED_RESOURCES = ("tpu", "nvidia.com/gpu", "cpu")
          SUPPORTS_CPU_MULTI_THREADING = True
  
          def __init__(self):
              super().__init__()
              self.device = xla_bridge.get_backend().platform
  
              self.model, self.state_dict = load_model(
                  bento_model, device=self.device
              )
              self.params = self.state_dict["params"]
              self.methods_cache: t.Dict[
                  str, t.Callable[..., t.Any]
              ] = {}
  
      def gen_run_method(self: FlaxRunnable, method_name: str):
          method = getattr(self.model, method_name)
          method_partial_kwargs = partial_kwargs.get(method_name)
          if method_partial_kwargs is not None:
              method = functools.partial(
                  method, **method_partial_kwargs
              )
  
          def mapping(
              item: jnp.ndarray | ext.NpNDArray | ext.PdDataFrame,
          ) -> jnp.ndarray:
              if LazyType["ext.NpNDArray"]("numpy.ndarray").isinstance(
                  item
              ):
                  return jnp.asarray(item)
              if LazyType["ext.PdDataFrame"](
                  "pandas.DataFrame"
              ).isinstance(item):
                  # NOTE: only to_numpy() are doing copying in memory here.
                  return jnp.asarray(item.to_numpy())
              return item
  
          def run_method(
              self: FlaxRunnable,
              *args: jnp.ndarray | ext.NpNDArray | ext.PdDataFrame,
          ):
              params = Params[jnp.ndarray](*args).map(mapping)
  
              arg = (
                  params.args[0]
                  if len(params.args) == 1
                  else params.args
              )
              # NOTE: can we jit this?
              # No?, as we should not interfere with JAX tracing in multiple threads
              # https://jax.readthedocs.io/en/latest/concurrency.html?highlight=concurrency
              return self.model.apply(
                  {"params": self.params}, arg, method=method
              )
  
          return run_method
  
      def add_runnable_method(
          method_name: str, options: ModelSignature
      ):
          def run_method(self: FlaxRunnable, *args: jnp.ndarray):
              fn = self.methods_cache.get(method_name)
              if not fn:
                  fn = gen_run_method(self, method_name)
                  self.methods_cache[method_name] = fn
              return fn(self, *args)
  
          FlaxRunnable.add_method(
              run_method,
              name=method_name,
              batchable=options.batchable,
              batch_dim=options.batch_dim,
              input_spec=options.input_spec,
              output_spec=options.output_spec,
          )
  
      for method_name, options in bento_model.info.signatures.items():
          add_runnable_method(method_name, options)
  
      return FlaxRunnable
  
imports: from __future__ import annotations
  import functools
  import logging
  import typing
  from pickle import UnpicklingError
  from types import ModuleType
  from typing import TYPE_CHECKING
  import msgpack.exceptions
  import bentoml
  from ...exceptions import BentoMLException
  from ...exceptions import MissingDependencyException
  from ...exceptions import NotFound
  from ..models.model import ModelContext
  from ..models.model import PartialKwargsModelOptions
  from ..runner.utils import Params
  from ..types import LazyType
  from ..utils import LazyLoader
  from ..utils.pkg import get_pkg_version
  from .common.jax import JaxArrayContainer
  from .common.jax import jax
  from .common.jax import jnp
label_desc: Use image node_id 4 for calling a global function with example usage: get_runnable(bento_model) and returns: FlaxRunnable, run_method, item, self, fn, jnp, jnp