code: def insert_zeros(
      message: ndarray, shape: Sequence[int], tail_capacity: int
  ) -> Tuple[ndarray, int]:
      """Insert zeros between head and tail based on tail capacity."""
      head_size = int(jnp.prod(jnp.array(shape)) * 4)
      splits = jnp.split(message, [head_size])
      needed_zeros = tail_capacity - splits[1].shape[0]
      tail_limit = needed_zeros + 1
  
      if needed_zeros < 0:
          raise ValueError(
              "tail_capacity insufficient for needed buffer size."
          )
  
      return (
          jnp.concatenate(
              (
                  splits[0],
                  jnp.zeros(needed_zeros, TAIL_DTYPE),
                  splits[1],
              )
          ),
          tail_limit,
      )
  
imports: from typing import NamedTuple, Sequence, Tuple
  import jax.lax
  import jax.numpy
  from jax.numpy import ndarray
  from jax.random import PRNGKey, randint
label_desc: Use image node_id 10 for calling a global function with example usage: insert_zeros(message, shape, tail_capacity) and returns: jnp,  tail_limit