code: def _extract_shapes(x):
      """
      Input:
          x: can be any of the following:
              - `nn.Module`
              - dict of shapes
              - dict of `InfShape`
              - str of path to a base shapes (.bsh) file
      Output:
          If `x` is dict of `InfShape`, then output itself.
          If `x` is path, then output a dict of `InfShapes` loaded from `x`.
          Else, output the shapes (not `InfShape`) associated to `x`
      """
      if isinstance(x, nn.Module):
          x_shapes = get_shapes(x)
      elif isinstance(x, dict):
          x_shapes = deepcopy(x)
      elif isinstance(x, str):
          # x is file name
          x_shapes = load_base_shapes(x)
      else:
          raise ValueError(f"unhandled x type: {type(x)}")
      return x_shapes
  
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 6 for calling a global function with example usage: _extract_shapes(x) and returns: x_shapes