code: def set_base_shapes(
      model,
      base,
      rescale_params=True,
      delta=None,
      savefile=None,
      do_assert=True,
  ):
      """Sets the `p.infshape` attribute for each parameter `p` of `model`.
  
      Inputs:
          model: nn.Module instance
          base: The base model.
              Can be nn.Module, a dict of shapes, a str, or None.
              If None, then defaults to `model`
              If str, then treated as filename for yaml encoding of a dict of base shapes.
          rescale_params:
              assuming the model is initialized using the default pytorch init (or
              He initialization etc that scale the same way with fanin): If True
              (default), rescales parameters to have the correct (μP) variances.
          do_assert:
      Output:
          same object as `model`, after setting the `infshape` attribute of each parameter.
      """
      if base is None:
          base = model
      base_shapes = _extract_shapes(base)
      if delta is not None:
          delta_shapes = _extract_shapes(delta)
          base_shapes = _zip_infshape_dict(base_shapes, delta_shapes)
      shapes = get_shapes(model)
      infshapes = _zip_infshape_dict(base_shapes, shapes)
      if savefile is not None:
          save_base_shapes(infshapes, savefile)
      apply_infshapes(model, infshapes)
      if do_assert:
          assert_hidden_size_inf(model)
      if rescale_params:
          for name, module in model.named_modules():
              if isinstance(module, MuReadout):
                  module._rescale_parameters()
              elif isinstance(module, (Linear, _ConvNd)):
                  rescale_linear_bias(module)
      return model
  
file_path: NeMo/nemo/collections/nlp/modules/common/megatron/mup/shape.py
imports: from copy import deepcopy
  import yaml
  from torch import nn
  from torch.nn import Linear
  from torch.nn.modules.conv import _ConvNd
  from nemo.collections.nlp.modules.common.megatron.mup.infshape import InfShape, zip_infshape
  from nemo.collections.nlp.modules.common.megatron.mup.layer import MuReadout, rescale_linear_bias
label_desc: Use image node_id 12 for calling a global function with example usage: set_base_shapes(model, base, rescale_params, delta, savefile, do_assert) and returns: model