code: def test_increment(K: int, include_weight: bool, 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")
      dtype_ = DTYPES[dtype]
  
      rms_layer = RMSNorm(K, include_weight=include_weight).cuda()
      x_orig = torch.rand(B, M, K, device=device, dtype=dtype_)
      y_orig = torch.rand(B, M, K, device=device, dtype=dtype_)
      x = x_orig.clone()
      y = y_orig.clone()
      if include_weight:
          torch.nn.init.normal_(rms_layer.weight)  # type: ignore
  
      context = torch.no_grad() if include_weight else nullcontext()
      with context:  # type: ignore
          baseline = rms_layer(x_orig + y_orig)
          out = rms_layer.increment_and_forward_(x, y)
      assert_allclose(out, baseline, atol=atol, rtol=rtol)
      assert_allclose(x, x_orig + y_orig, atol=atol, rtol=rtol)
      assert_allclose(y, y_orig, 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 2 for calling a global function with example usage: test_increment(K, include_weight, dtype) without return types