code: def apply_chunking_to_forward(
      forward_fn: Callable[..., torch.Tensor],
      chunk_size: int,
      chunk_dim: int,
      *input_tensors,
  ) -> torch.Tensor:
      """
      This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
      `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
  
      If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
      applying `forward_fn` to `input_tensors`.
  
      Args:
          forward_fn (`Callable[..., torch.Tensor]`):
              The forward function of the model.
          chunk_size (`int`):
              The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
          chunk_dim (`int`):
              The dimension over which the `input_tensors` should be chunked.
          input_tensors (`Tuple[torch.Tensor]`):
              The input tensors of `forward_fn` which will be chunked
  
      Returns:
          `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
  
  
      Examples:
  
      ```python
      # rename the usual forward() fn to forward_chunk()
      def forward_chunk(self, hidden_states):
          hidden_states = self.decoder(hidden_states)
          return hidden_states
  
  
      # implement a chunked forward function
      def forward(self, hidden_states):
          return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
      ```"""
  
      assert (
          len(input_tensors) > 0
      ), f"{input_tensors} has to be a tuple/list of tensors"
  
      # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
      num_args_in_forward_chunk_fn = len(
          inspect.signature(forward_fn).parameters
      )
      if num_args_in_forward_chunk_fn != len(input_tensors):
          raise ValueError(
              f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
              "tensors are given"
          )
  
      if chunk_size > 0:
          tensor_shape = input_tensors[0].shape[chunk_dim]
          for input_tensor in input_tensors:
              if input_tensor.shape[chunk_dim] != tensor_shape:
                  raise ValueError(
                      f"All input tenors have to be of the same shape: {tensor_shape}, "
                      f"found shape {input_tensor.shape[chunk_dim]}"
                  )
  
          if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
              raise ValueError(
                  f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                  f"size {chunk_size}"
              )
  
          num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size
  
          # chunk input tensor into tuples
          input_tensors_chunks = tuple(
              input_tensor.chunk(num_chunks, dim=chunk_dim)
              for input_tensor in input_tensors
          )
          # apply forward fn to every tuple
          output_chunks = tuple(
              forward_fn(*input_tensors_chunk)
              for input_tensors_chunk in zip(*input_tensors_chunks)
          )
          # concatenate output at same dimension
          return torch.cat(output_chunks, dim=chunk_dim)
  
      return forward_fn(*input_tensors)
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/pytorch_utils.py
imports: import inspect
  from typing import Callable, List, Optional, Set, Tuple, Union
  import torch
  from packaging import version
  from torch import _softmax_backward_data, nn
  from .utils import logging
label_desc: Use image node_id 6 for calling a global function with example usage: apply_chunking_to_forward(forward_fn, chunk_size, chunk_dim) and returns: forward_fn, torch