class_docstr: Plugin for training with fp8 precision via nvidia's
  `Transformer Engine <https://docs.nvidia.com/deeplearning/transformer-engine>`__.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  Args:
      weights_dtype: The weights dtype to use.
      recipe: Recipe for the DelayedScaling
          `configuration <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling>`__.
          In dict format or the dataclass format.
      replace_layers: Whether to replace ``Linear`` and ``LayerNorm`` layers automatically with their Transformer
          Engine alternatives. Note that they don't subclass the torch equivalents so checks like
          ``isinstance(l, torch.nn.Linear)`` will not pass.
      fallback_compute_dtype: The compute dtype to use for operations that don't support fp8 autocast. Defaults to the
          same as ``weights_dtype``.
  
  .. note::
  
      Support for FP8 in the linear layers with this plugin is currently limited to tensors
      with shapes where the dimensions are divisible by 8 and 16 respectively. You might want to add padding to your
      inputs to conform to this restriction.
class_docstr_tok: ["Plugin","for","training","with","fp8","precision","via","nvidia's","`","Transformer","Engine","<","https",":","\/\/docs.nvidia.com\/deeplearning\/transformer-engine",">","`","__",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","Args",":","weights_dtype",":","The","weights","dtype","to","use",".","recipe",":","Recipe","for","the","DelayedScaling","`","configuration","<","https",":","\/\/docs.nvidia.com\/deeplearning\/transformer-engine\/user-guide\/api\/common.html","#","transformer_engine.common.recipe.DelayedScaling",">","`","__",".","In","dict","format","or","the","dataclass","format",".","replace_layers",":","Whether","to","replace","``","Linear","``","and","``","LayerNorm","``","layers","automatically","with","their","Transformer","Engine","alternatives",".","Note","that","they","do","n't","subclass","the","torch","equivalents","so","checks","like","``","isinstance","(","l",",","torch.nn.Linear",")","``","will","not","pass",".","fallback_compute_dtype",":","The","compute","dtype","to","use","for","operations","that","do","n't","support","fp8","autocast",".","Defaults","to","the","same","as","``","weights_dtype","``",".","..","note",":",":","Support","for","FP8","in","the","linear","layers","with","this","plugin","is","currently","limited","to","tensors","with","shapes","where","the","dimensions","are","divisible","by","8","and","16","respectively",".","You","might","want","to","add","padding","to","your","inputs","to","conform","to","this","restriction","."]
code: def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
      # avoid converting if any is found. assume the user took care of it
      if any(
          "transformer_engine.pytorch" in m.__module__
          for m in module.modules()
      ):
          if self.replace_layers is True:
              # info level because this is expected with `init_module`
              rank_zero_info(
                  "`TransformerEnginePrecision(replace_layers=True)` is set but the model already contains"
                  " TransformerEngine layers. Skipping"
              )
      elif self.replace_layers in (None, True):
          _convert_layers(module)
      module = module.to(dtype=self.weights_dtype)
      return module
  
imports: import logging
  from contextlib import ExitStack
  from typing import TYPE_CHECKING, Any, ContextManager, Literal, Mapping, Optional, Union
  import torch
  from lightning_utilities import apply_to_collection
  from lightning_utilities.core.imports import RequirementCache
  from torch import Tensor
  from typing_extensions import override
  from lightning.fabric.plugins.precision.precision import Precision
  from lightning.fabric.plugins.precision.utils import _ClassReplacementContextManager, _convert_fp_tensor, _DtypeContextManager
  from lightning.fabric.utilities.rank_zero import rank_zero_info, rank_zero_warn
label_desc: Use image node_id 2 for calling the TransformerEnginePrecision obj's underlying member method code with example usage: obj.convert_module(module) and returns: module