class_docstr: Generates a summary of all layers in a :class:`~lightning.pytorch.core.LightningModule`.
  
  Args:
      model: The model to summarize (also referred to as the root module).
  
      max_depth: Maximum depth of modules to show. Use -1 to show all modules or 0 to show no
          summary. Defaults to 1.
  
  The string representation of this summary prints a table with columns containing
  the name, type and number of parameters for each layer.
  
  The root module may also have an attribute ``example_input_array`` as shown in the example below.
  If present, the root module will be called with it as input to determine the
  intermediate input- and output shapes of all layers. Supported are tensors and
  nested lists and tuples of tensors. All other types of inputs will be skipped and show as `?`
  in the summary table. The summary will also display `?` for layers not used in the forward pass.
  If there are parameters not associated with any layers or modules, the count of those parameters
  will be displayed in the table under `other params`. The summary will display `n/a` for module type,
  in size, and out size.
  
  Example::
  
      >>> import lightning.pytorch as pl
      >>> class LitModel(pl.LightningModule):
      ...
      ...     def __init__(self):
      ...         super().__init__()
      ...         self.net = nn.Sequential(nn.Linear(256, 512), nn.BatchNorm1d(512))
      ...         self.example_input_array = torch.zeros(10, 256)  # optional
      ...
      ...     def forward(self, x):
      ...         return self.net(x)
      ...
      >>> model = LitModel()
      >>> ModelSummary(model, max_depth=1)  # doctest: +NORMALIZE_WHITESPACE
        | Name | Type       | Params | In sizes  | Out sizes
      ------------------------------------------------------------
      0 | net  | Sequential | 132 K  | [10, 256] | [10, 512]
      ------------------------------------------------------------
      132 K     Trainable params
      0         Non-trainable params
      132 K     Total params
      0.530     Total estimated model params size (MB)
      >>> ModelSummary(model, max_depth=-1)  # doctest: +NORMALIZE_WHITESPACE
        | Name  | Type        | Params | In sizes  | Out sizes
      --------------------------------------------------------------
      0 | net   | Sequential  | 132 K  | [10, 256] | [10, 512]
      1 | net.0 | Linear      | 131 K  | [10, 256] | [10, 512]
      2 | net.1 | BatchNorm1d | 1.0 K    | [10, 512] | [10, 512]
      --------------------------------------------------------------
      132 K     Trainable params
      0         Non-trainable params
      132 K     Total params
      0.530     Total estimated model params size (MB)
class_docstr_tok: ["Generates","a","summary","of","all","layers","in","a",":","class",":","`","~lightning.pytorch.core.LightningModule","`",".","Args",":","model",":","The","model","to","summarize","(","also","referred","to","as","the","root","module",")",".","max_depth",":","Maximum","depth","of","modules","to","show",".","Use","-1","to","show","all","modules","or","0","to","show","no","summary",".","Defaults","to","1",".","The","string","representation","of","this","summary","prints","a","table","with","columns","containing","the","name",",","type","and","number","of","parameters","for","each","layer",".","The","root","module","may","also","have","an","attribute","``","example_input_array","``","as","shown","in","the","example","below",".","If","present",",","the","root","module","will","be","called","with","it","as","input","to","determine","the","intermediate","input-","and","output","shapes","of","all","layers",".","Supported","are","tensors","and","nested","lists","and","tuples","of","tensors",".","All","other","types","of","inputs","will","be","skipped","and","show","as","`","?","`","in","the","summary","table",".","The","summary","will","also","display","`","?","`","for","layers","not","used","in","the","forward","pass",".","If","there","are","parameters","not","associated","with","any","layers","or","modules",",","the","count","of","those","parameters","will","be","displayed","in","the","table","under","`","other","params","`",".","The","summary","will","display","`","n\/a","`","for","module","type",",","in","size",",","and","out","size",".","Example",":",":",">",">",">","import","lightning.pytorch","as","pl",">",">",">","class","LitModel","(","pl.LightningModule",")",":","...","...","def","__init__","(","self",")",":","...","super","(",")",".__init__","(",")","...","self.net","=","nn.Sequential","(","nn.Linear","(","256",",","512",")",",","nn.BatchNorm1d","(","512",")",")","...","self.example_input_array","=","torch.zeros","(","10",",","256",")","#","optional","...","...","def","forward","(","self",",","x",")",":","...","return","self.net","(","x",")","...",">",">",">","model","=","LitModel","(",")",">",">",">","ModelSummary","(","model",",","max_depth=1",")","#","doctest",":","+NORMALIZE_WHITESPACE","|","Name","|","Type","|","Params","|","In","sizes","|","Out","sizes","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","0","|","net","|","Sequential","|","132","K","|","[","10",",","256","]","|","[","10",",","512","]","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","132","K","Trainable","params","0","Non-trainable","params","132","K","Total","params","0.530","Total","estimated","model","params","size","(","MB",")",">",">",">","ModelSummary","(","model",",","max_depth=-1",")","#","doctest",":","+NORMALIZE_WHITESPACE","|","Name","|","Type","|","Params","|","In","sizes","|","Out","sizes","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","0","|","net","|","Sequential","|","132","K","|","[","10",",","256","]","|","[","10",",","512","]","1","|","net.0","|","Linear","|","131","K","|","[","10",",","256","]","|","[","10",",","512","]","2","|","net.1","|","BatchNorm1d","|","1.0","K","|","[","10",",","512","]","|","[","10",",","512","]","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","132","K","Trainable","params","0","Non-trainable","params","132","K","Total","params","0.530","Total","estimated","model","params","size","(","MB",")"]
code: def __init__(
      self, model: "pl.LightningModule", max_depth: int = 1
  ) -> None:
      self._model = model
  
      if not isinstance(max_depth, int) or max_depth < -1:
          raise ValueError(
              f"`max_depth` can be -1, 0 or > 0, got {max_depth}."
          )
  
      self._max_depth = max_depth
      self._layer_summary = self.summarize()
      # 1 byte -> 8 bits
      # TODO: how do we compute precision_megabytes in case of mixed precision?
      precision_to_bits = {"64": 64, "32": 32, "16": 16, "bf16": 16}
      precision = (
          precision_to_bits.get(self._model.trainer.precision, 32)
          if self._model._trainer
          else 32
      )
      self._precision_megabytes = (precision / 8.0) * 1e-6
  
imports: import contextlib
  import logging
  import math
  from collections import OrderedDict
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
  from torch.utils.hooks import RemovableHandle
  import lightning.pytorch
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch.utilities.model_helpers import _ModuleMode
  from lightning.pytorch.utilities.rank_zero import WarningCache
label_desc: Use image node_id 1 to create a new ModelSummary object  with example: obj = ModelSummary(model, max_depth)