code: def assert_allclose(
      # The output of the tested function
      out: torch.Tensor,
      # The output of the reference implementation
      ref: torch.Tensor,
      # The output of the reference implementation in f32
      ref32: Optional[torch.Tensor] = None,
      msg: str = "failed",
      atol: Optional[float] = None,
      rtol: Optional[float] = None,
  ) -> None:
      """
      Improved version of
      ```
      assert torch.allclose(out, ref)
      ```
  
      Except that we provide useful error message, and also compare
      to the output of the f32 calculation.
      """
      out = out.float()
      ref = ref.float()
      if atol is None:
          atol = 1e-8
      if rtol is None:
          rtol = 1e-5
      assert out.shape == ref.shape
      compare_to = ref32 if ref32 is not None else ref
      assert out.shape == compare_to.shape
      if torch.allclose(out, ref, rtol=rtol, atol=atol) or (
          ref32 is not None
          and torch.allclose(out, ref32, rtol=rtol, atol=atol)
      ):
          return
  
      flatten_diff = (
          (out - compare_to).abs() - atol - compare_to.abs() * rtol
      ).flatten()
      max_pos = flatten_diff.argmax()
  
      if ref32 is not None:
          flatten_diff_vsf32 = (
              (ref - ref32).abs() - atol - ref32.abs() * rtol
          ).flatten()
          max_pos_vsf32 = flatten_diff_vsf32.argmax()
          assert False, (
              f"{msg}: "
              f"out={out.flatten()[max_pos]} and ref32={compare_to.flatten()[max_pos]} (diff={flatten_diff[max_pos]} > 0)"
              f"/ atol={atol}, rtol={rtol}.\n"
              f"NOTE: ref vs ref_f32:\n"
              f"ref={ref.flatten()[max_pos_vsf32]} and ref32={ref32.flatten()[max_pos_vsf32]} "
              f"(diff={flatten_diff_vsf32[max_pos_vsf32]})"
          )
      else:
          assert False, (
              f"{msg}: "
              f"out={out.flatten()[max_pos]} and ref={compare_to.flatten()[max_pos]} (diff={flatten_diff[max_pos]} > 0)"
              f"/ atol={atol}, rtol={rtol}"
          )
  
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 1 for calling a global function with example usage: assert_allclose(out, ref, ref32, msg, atol, rtol) without return types