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 __init__(
      self,
      *,
      weights_dtype: torch.dtype,
      recipe: Optional[
          Union[Mapping[str, Any], "DelayedScaling"]
      ] = None,
      replace_layers: Optional[bool] = None,
      fallback_compute_dtype: Optional[torch.dtype] = None,
  ) -> None:
      if not _TRANSFORMER_ENGINE_AVAILABLE:
          raise ModuleNotFoundError(str(_TRANSFORMER_ENGINE_AVAILABLE))
      from transformer_engine.common.recipe import DelayedScaling
  
      if recipe is None:
          recipe = DelayedScaling()
      elif isinstance(recipe, Mapping):
          recipe = dict(recipe)  # copy
          if "fp8_format" in recipe:
              from transformer_engine.common.recipe import Format
  
              recipe["fp8_format"] = getattr(
                  Format, recipe["fp8_format"]
              )
          recipe = DelayedScaling(**recipe)
  
      self.weights_dtype = weights_dtype
      self.recipe = recipe
      self.replace_layers = replace_layers
      self.fallback_compute_dtype = (
          fallback_compute_dtype or weights_dtype
      )
  
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 1 to create a new TransformerEnginePrecision object from inherited base classes: Precision with example: obj = TransformerEnginePrecision()