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
Private API: use :obj:`~bentoml.Model.to_runnable` instead.
src/bentoml/_internal/frameworks/flax.py