class_docstr: Handle a layer normalization, like torch.nn.LayerNorm_.
  
  This implementation should be measurably faster than the default PyTorch layernorm (as of PyTorch 1.9),
  both for training and inference worloads.
  
  .. NOTE: Computations under Torch AMP are kept as float32 by default, one can change this to be float16
      by setting the flag `xformers.triton.k_layer_norm._triton_layernorm_fp16_enabled = True`
  
  .. _torch.nn.LayerNorm: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html
class_docstr_tok: ["Handle","a","layer","normalization",",","like","torch.nn.LayerNorm_",".","This","implementation","should","be","measurably","faster","than","the","default","PyTorch","layernorm","(","as","of","PyTorch","1.9",")",",","both","for","training","and","inference","worloads",".","..","NOTE",":","Computations","under","Torch","AMP","are","kept","as","float32","by","default",",","one","can","change","this","to","be","float16","by","setting","the","flag","`","xformers.triton.k_layer_norm._triton_layernorm_fp16_enabled","=","True","`","..","_torch.nn.LayerNorm",":","https",":","\/\/pytorch.org\/docs\/stable\/generated\/torch.nn.LayerNorm.html"]
code: def __init__(self, normalized_shape, affine=True, eps=1e-06):
      super().__init__()
      if affine:
          self.weight = nn.Parameter(torch.ones(normalized_shape))
          self.bias = nn.Parameter(torch.zeros(normalized_shape))
      else:
          self.weight = self.bias = None
      self.epsilon = eps
  
file_path: xformers/xformers/triton/layer_norm.py
imports: import logging
  from typing import Optional
  import torch
  import torch.nn
  import triton
  from torch.cuda.amp import custom_bwd, custom_fwd
  from xformers.triton.k_layer_norm import layer_norm_bwd_dwdb, layer_norm_bwd_dx_fused, layer_norm_fw
label_desc: Use image node_id 1 to create a new FusedLayerNorm object from inherited base classes: nn with example: obj = FusedLayerNorm(normalized_shape, affine, eps)