code: def optimal_step_size(
      last_step,
      mean_error_ratio,
      safety=0.9,
      ifactor=10.0,
      dfactor=0.2,
      order=5.0,
  ):
      """Compute optimal Runge-Kutta stepsize."""
      dfactor = jnp.where(mean_error_ratio < 1, 1.0, dfactor)
  
      factor = jnp.minimum(
          ifactor,
          jnp.maximum(
              mean_error_ratio ** (-1.0 / order) * safety, dfactor
          ),
      )
      return jnp.where(
          mean_error_ratio == 0, last_step * ifactor, last_step * factor
      )
  
imports: from functools import partial
  import operator
  import jax
  import jax.numpy
  from jax._src import core
  from jax import custom_derivatives
  from jax import lax
  from jax._src.numpy.util import promote_dtypes_inexact
  from jax._src.util import safe_map, safe_zip
  from jax.flatten_util import ravel_pytree
  from jax.tree_util import tree_leaves, tree_map
  from jax._src import linear_util
label_desc: Use image node_id 9 for calling a global function with example usage: optimal_step_size(last_step, mean_error_ratio, safety, ifactor, dfactor, order) and returns: jnp