code: def prune_module_summary(
      module_summary: ModuleSummary, *, max_depth: int
  ) -> None:
      """
      Prune the module summaries that are deeper than max_depth in the module
      summary tree. The ModuleSummary object is prunned inplace.
  
      Args:
          module_summary: Root module summary to prune.
          max_depth: The maximum depth of module summaries to keep.
  
      Raises:
          ValueError:
              If `max_depth` is an int less than 1
      """
      if max_depth < 1:
          raise ValueError(
              f"`max_depth` must be an int greater than 0. Got {max_depth}."
          )
      if max_depth == 1:
          module_summary._submodule_summaries = {}
          return
  
      for (
          submodule_summary
      ) in module_summary._submodule_summaries.values():
          prune_module_summary(
              submodule_summary, max_depth=max_depth - 1
          )
  
imports: import copy
  import math
  import warnings
  from collections import defaultdict, deque
  from dataclasses import dataclass, field
  from enum import Enum
  from time import perf_counter
  from typing import Any, Callable, DefaultDict, Deque, Dict, List, MutableMapping, Optional, Tuple, Union
  import torch
  from torch.nn.parameter import UninitializedParameter
  from torch.utils._pytree import PyTree, tree_flatten
  from torch.utils.hooks import RemovableHandle
  from torchtnt.utils.version import is_torch_version_geq_1_13
  from typing_extensions import Literal
label_desc: Use image node_id 8 for calling a global function with example usage: prune_module_summary(module_summary) without return types