code: def bind_rng_to_host_device(
      rng: jnp.ndarray, axis_name: str, bind_to: Optional[str] = None
  ) -> jnp.ndarray:
      """Bind rng to host device."""
      if bind_to is None:
          return rng
      if bind_to == "host":
          return jax.random.fold_in(rng, jax.process_index())
      elif bind_to == "device":
          return jax.random.fold_in(rng, jax.lax.axis_index(axis_name))
      else:
          raise ValueError(
              "`bind_to` should be one of the `[None, 'host', 'device']`"
          )
  
imports: from typing import Tuple, Optional
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: bind_rng_to_host_device(rng, axis_name, bind_to) and returns: rng, jax, jax