code: def test_fused_matmul(shape, dtype):
      """Check that the matrix multiply kernel and Pytorch's give the same results"""
      # TODO: fix or remove this
      pytest.skip("This is broken")
      torch.random.manual_seed(0)
  
      # Raw fused matrix multiply first, to catch gross errors
      a = torch.normal(
          0, 1, size=(shape[-2], shape[-1]), dtype=dtype, device="cuda"
      )
      b = torch.normal(
          0, 1, size=(shape[-1], shape[-2]), dtype=dtype, device="cuda"
      )
  
      # Test that not passing any bias is fine
      res_torch = a @ b
      res_triton, _ = fused_matmul(
          a, b.transpose(0, 1).contiguous(), bias=None, activation=0
      )
      torch.testing.assert_close(res_torch, res_triton)
  
      # Now test with a real FMA
      c = -torch.randn((shape[-2],), dtype=dtype, device="cuda")
      res_torch = torch.addmm(c, a, b)
      res_triton, _ = fused_matmul(a, b.transpose(1, 0).contiguous(), c)
  
      torch.testing.assert_close(
          res_torch,
          res_triton,
          atol=1e-3,
          rtol=1e-3,
          msg="Fused matmul broken",
      )
  
      # Now check that adding an activation to the mix still produces valid results
      # NOTE: SquaredReLU fails, some outlier representation issue but the eyeballed results look reasonable
      # could be due to a different accumulation out of the box (tf32 for instance)
      for activation in filter(
          lambda x: x
          not in (Activation.SquaredReLU, Activation.StarReLU),
          Activation,
      ):
          torch_activation = build_activation(activation.value)
          res_torch = torch_activation(torch.addmm(c, a, b))
  
          triton_activation_index = get_triton_activation_index(
              activation
          )
          print(activation, triton_activation_index)
          res_triton, _ = fused_matmul(
              a,
              b.transpose(1, 0).contiguous(),
              c,
              triton_activation_index,
          )
  
          torch.testing.assert_close(
              res_torch,
              res_triton,
              atol=1e-3,
              rtol=1e-3,
              msg=f"Fused matmul broken with activation {activation}",
          )
  
file_path: xformers/tests/test_triton_fused_linear.py
imports: import logging
  import pytest
  import torch
  from torch.cuda.amp.autocast_mode import autocast
  import xformers
  from xformers.components import Activation, build_activation
label_desc: Use image node_id 1 for calling a global function with example usage: test_fused_matmul(shape, dtype) without return types