code: def stack_module_state(
      models: List[nn.Module],
  ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
      """stack_module_state(models) -> params, buffers
  
      Prepares a list of torch.nn.Modules for ensembling with :func:`vmap`.
  
      Given a list of ``M`` ``nn.Modules`` of the same class, returns two dictionaries
      that stack all of their parameters and buffers together, indexed by name.
      The stacked parameters are optimizable (i.e. they are new leaf nodes in the
      autograd history that are unrelated to the original parameters and can be
      passed directly to an optimizer).
  
      Here's an example of how to ensemble over a very simple model:
  
      .. code-block:: python
  
          num_models = 5
          batch_size = 64
          in_features, out_features = 3, 3
          models = [torch.nn.Linear(in_features, out_features) for i in range(num_models)]
          data = torch.randn(batch_size, 3)
  
          def wrapper(params, buffers, data):
              return torch.func.functional_call(model[0], (params, buffers), data)
  
          params, buffers = stack_module_state(models)
          output = vmap(wrapper, (0, 0, None))(params, buffers, data)
  
          assert output.shape == (num_models, batch_size, out_features)
  
      When there's submodules, this follows state dict naming conventions
  
      .. code-block:: python
  
          import torch.nn as nn
          class Foo(nn.Module):
              def __init__(self, in_features, out_features):
                  super().__init__()
                  hidden = 4
                  self.l1 = nn.Linear(in_features, hidden)
                  self.l2 = nn.Linear(hidden, out_features)
  
              def forward(self, x):
                  return self.l2(self.l1(x))
  
          num_models = 5
          in_features, out_features = 3, 3
          models = [Foo(in_features, out_features) for i in range(num_models)]
          params, buffers = stack_module_state(models)
          print(list(params.keys()))  # "l1.weight", "l1.bias", "l2.weight", "l2.bias"
  
      .. warning::
          All of the modules being stacked together must be the same (except for
          the values of their parameters/buffers). For example, they should be in the
          same mode (training vs eval).
      """
      if len(models) == 0:
          raise RuntimeError(
              "stack_module_state: Expected at least one model, got 0."
          )
      if not (
          all(m.training for m in models)
          or all(not m.training for m in models)
      ):
          raise RuntimeError(
              "stack_module_state: Expected all models to have the same training/eval mode."
          )
      model0_typ = type(models[0])
      if not all(type(m) == model0_typ for m in models):
          raise RuntimeError(
              "stack_module_state: Expected all models to be of the same class."
          )
      all_params = [dict(model.named_parameters()) for model in models]
      params = {
          k: construct_stacked_leaf(
              tuple(params[k] for params in all_params), k
          )
          for k in all_params[0]
      }
      all_buffers = [dict(model.named_buffers()) for model in models]
      buffers = {
          k: construct_stacked_leaf(
              tuple(buffers[k] for buffers in all_buffers), k
          )
          for k in all_buffers[0]
      }
  
      return params, buffers
  
file_path: pytorch/torch/_functorch/functional_call.py
imports: from collections import Counter
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
  from torch._functorch.utils import exposed_in
label_desc: Use image node_id 2 for calling a global function with example usage: stack_module_state(models) and returns: params,  buffers