code: def test_forward(K: int, dtype: str):
      atol = 1e-8 if dtype == "f32" else 1e-4
      rtol = 1e-5 if dtype == "f32" else 0.01
      torch.manual_seed(1)
      B, M, K = 31, 27, K
      device = torch.device("cuda")
  
      rms_layer = RMSNorm(K).cuda()
      baseline_layer = RMSNormPytorch(K).cuda()
      x = torch.rand(B, M, K, device=device, dtype=DTYPES[dtype])
      torch.nn.init.normal_(rms_layer.weight)  # type: ignore
      with torch.no_grad():
          x_rms = rms_layer(x)
          assert x_rms.shape == x.shape
          baseline_layer.weight.copy_(rms_layer.weight)  # type: ignore
      baseline = baseline_layer(x)
      assert_allclose(x_rms, baseline, atol=atol, rtol=rtol)
  
      torch.nn.init.ones_(rms_layer.weight)  # type: ignore
      with torch.no_grad():
          x_rms1 = rms_layer(x)
      assert not torch.allclose(x_rms, x_rms1)
      rms1_layer = RMSNorm(K, include_weight=False)
      with torch.no_grad():
          x_rms_1 = rms1_layer(x)
      assert_allclose(x_rms1, x_rms_1, atol=atol, rtol=rtol)
  
file_path: xformers/tests/test_rmsnorm.py
imports: from contextlib import nullcontext
  from typing import Optional
  import pytest
  import torch
  from torch import nn
  from xformers.ops import RMSNorm
  from .utils import assert_allclose
label_desc: Use image node_id 1 for calling a global function with example usage: test_forward(K, dtype) without return types