# Copied from https://github.com/google-research/big_vision

                    static_return=False):

  """Wraps a function with code that pads, shards, then un-shards, un-pads.

  Args:

    wrapped: the function to be wrapped. Signature is `params, *args, *kwargs`.

    static_argnums: indices of arguments to `wrapped` that should _not_ be

      padded and sharded, but instead be forwarded as-is. The default is (0,)

      because by far the most common use-case is to pass `params` first.

    static_argnames: names of kwargs to `wrapped` that should _not_ be padded

      and sharded, but instead be forwarded as-is.

    static_return: whether not to un-shard, and un-pad the return value; static

      return values are typically used with eval steps that compute metrics

  Returns:

    A new function that pads and shards its arguments before passing them to

    the wrapped function, and un-shards and un-pads the returned pytree.

    This is useful for calling a pmap'ed function with inputs that aren't

    divisible by the number of devices. A typical use is:

      @pad_shard_unpad

      @jax.pmap

      def forward(params, x): ...

  Notes:

    The padding is done in host-memory before being passed to the function, and

    the values returned by the function are transferred back to host memory.

    The returned function is augmented with a new keyword-only argument

    `min_device_batch` that, if specified, forces padding inputs to at least

    this size per device. This can be useful to avoid recompiles for the last

    batch and reduce memory fragmentation.

    For more information refer to

    https://flax.readthedocs.io/en/latest/howtos/full_eval.html

  """

            [x, np.zeros((d * (min_device_batch - db), *shape), x.dtype)])

      # Transfer back before cutting, to reduce on-device shape diversity.