code: def test_forward_backward(
      shape,
      device,
      op,
      dtype,
      autocast: bool,
      pack_weights: bool,
      bias: bool,
  ):
      torch.manual_seed(shape[0] * shape[1] * shape[2])
      FORWARD_ATOL = {
          torch.float: 2e-6,
          torch.half: 1e-2,
          torch.bfloat16: 1e-2,
      }
      FORWARD_RTOL = {
          torch.float: 1e-5,
          torch.half: 4e-3,
          torch.bfloat16: 4e-3,
      }
      BACKWARD_ATOL = {
          torch.float: 3e-4,
          torch.half: 0.5,
          torch.bfloat16: 4.0,  # !!
      }
      BACKWARD_RTOL = {
          torch.float: 2e-3,
          torch.half: 1e-2,
          torch.bfloat16: 4e-2,
      }
  
      if not op.supports(
          xsw.SwiGLUOpDispatch(
              device=device,
              dtype=dtype,
              dtype_autocast_gpu=dtype
              if autocast and device == "cuda"
              else None,
              packed_weights=pack_weights,
              bias_enabled=bias,
          )
      ):
          pytest.skip("Not supported by operator")
  
      inp_model_dtype = torch.float if autocast else dtype
      x = torch.randn(shape[:2], device=device, dtype=inp_model_dtype)
  
      module = copy.deepcopy(
          create_module_cached(
              in_features=shape[1],
              hidden_features=shape[2],
              bias=bias,
              _pack_weights=pack_weights,
          )
      )
      x_f32: Optional[torch.Tensor]
      ref_f32: Optional[torch.Tensor]
      module_f32: Optional[torch.nn.Module]
      if dtype != torch.float:
          x_f32, module_f32 = x.to(device).to(torch.float), module.to(
              device
          )
          x_f32.requires_grad_()
          ref_f32 = module_f32(x_f32)
      else:
          x_f32, module_f32, ref_f32 = None, None, None
  
      x, module = x.to(device).to(inp_model_dtype), module.to(
          device
      ).to(inp_model_dtype)
      x.requires_grad_()
  
      # Forward
      cm = cast(
          ContextManager,
          torch.autocast("cuda", dtype=dtype)
          if autocast
          else nullcontext(),
      )
      with cm:
          ref = module(x)
          out = xsw.swiglu(x, *module._ordered_params(), op=op)
  
      if ref_f32 is None:
          ref_f32 = ref
  
      assert_allclose(
          out,
          ref,
          ref_f32,
          "fw",
          atol=FORWARD_ATOL[dtype],
          rtol=FORWARD_RTOL[dtype],
      )
  
      # Backward
      grad = torch.randn_like(ref)
  
      def backward_gather_grads(inp, output):
          output.backward(grad.to(output.dtype))
          grads = {}
          for name, param in module.named_parameters():
              grads[name] = param.grad.clone()
              param.grad = None
          grads["x"] = inp.grad.clone()
          inp.grad = None
          return grads
  
      grads_ref = backward_gather_grads(x, ref)
      grads_out = backward_gather_grads(x, out)
      grads_ref32 = (
          backward_gather_grads(x_f32, ref_f32)
          if module_f32 is not None
          else grads_ref
      )
  
      assert list(grads_ref.keys()) == list(grads_out.keys())
      for name, gref in grads_ref.items():
          gout = grads_out[name]
          assert_allclose(
              gout,
              gref,
              grads_ref32.get(name),
              f"{name}.grad",
              atol=BACKWARD_ATOL[dtype],
              rtol=BACKWARD_RTOL[dtype],
          )
          # Ensure `gout >> atol`, so that the test is meaningful
          assert (
              gout.norm(2) > BACKWARD_ATOL[dtype] / BACKWARD_RTOL[dtype]
          )
  
file_path: xformers/tests/test_swiglu.py
imports: import copy
  import functools
  import random
  from contextlib import nullcontext
  from typing import ContextManager, Optional, Sequence, cast
  import pytest
  import torch
  import xformers.ops.swiglu_op
label_desc: Use image node_id 4 for calling a global function with example usage: test_forward_backward(shape, device, op, dtype, autocast, pack_weights, bias) and returns: grads