code: def test_layernorm_parity(shape, amp):
      """Check that PyTorch and Triton softmax give the same result"""
  
      # Get the same inputs
      torch.random.manual_seed(0)
      X = torch.normal(
          0, 1, size=shape, device="cuda", requires_grad=True
      )
  
      torch.random.manual_seed(0)
      X_ = torch.normal(
          0, 1, size=shape, device="cuda", requires_grad=True
      )
  
      eps = 1e-4
  
      # Initialize the two layers, weights are 1 and 0 by default, no randomness
      torch_layernorm = torch.nn.LayerNorm(X.shape[-1], eps=eps).to(
          "cuda"
      )
      triton_layernorm = FusedLayerNorm(
          X.shape[-1], affine=True, eps=eps
      ).to("cuda")
  
      with autocast(enabled=amp):
          assert torch.allclose(
              X, X_
          )  # sanity checking, else all hell breaks loose
  
          # Check the forward pass
          y_torch = torch_layernorm(X)
          y_triton = triton_layernorm(X_)
          assert torch.allclose(
              y_torch.norm(), y_triton.norm(), atol=1e-3
          ), f"{torch.norm(y_torch)} vs. {torch.norm(y_triton)}"
  
          # Check that BW also gives the same result
          loss_torch = torch.norm(y_torch)
          loss_torch.backward()
  
          loss_triton = torch.norm(y_triton)
          loss_triton.backward()
  
          print(torch.norm(y_torch), torch.norm(y_triton))
  
          print(y_torch[0, :])
          print(y_triton[0, :])
  
          # There are 3 items to check:
          # - gradient on the inputs
          assert torch.allclose(
              X.grad, X_.grad
          ), f"Inputs grad mismatch: {torch.norm(X.grad)} vs. {torch.norm(X_.grad)}"
  
          # - gradient on the layernorm weight
          assert torch.allclose(
              torch_layernorm.weight.grad,
              triton_layernorm.weight.grad,
              atol=1e-3,
          ), (
              f"Weight grad mismatch: {torch.norm(torch_layernorm.weight.grad)} vs."
              + f" {torch.norm(triton_layernorm.weight.grad)}"
          )
  
          # - gradient on the layernorm bias
          assert torch.allclose(
              torch_layernorm.bias.grad,
              triton_layernorm.bias.grad,
              atol=1e-3,
          ), (
              f"Bias grad mismatch: {torch.norm(torch_layernorm.bias.grad)} vs."
              + f" {torch.norm(triton_layernorm.bias.grad)}"
          )
  
imports: import logging
  import pytest
  import torch
  from torch.cuda.amp.autocast_mode import autocast
  import xformers
label_desc: Use image node_id 1 for calling a global function with example usage: test_layernorm_parity(shape, amp) without return types