code: def test_bmm(dtype, batch_size, a_shape, b_shape, b_axis, device):
      if dtype == torch.float16 and device.type == "cpu":
          pytest.skip(
              "Matrix multiplication is not supported for float16 on CPU."
          )
      qa = random_qtensor((batch_size,) + a_shape, dtype=dtype).to(
          device
      )
      qb = random_qtensor(
          (batch_size,) + b_shape, axis=b_axis, dtype=dtype
      ).to(device)
      qbmm = torch.bmm(qa, qb)
      assert isinstance(qbmm, QTensor)
      # The outputs should be almost identical if we use the dequantized inputs
      bmm = torch.bmm(qa.dequantize(), qb.dequantize())
      # We need to increase atol and rtol for float16
      atol = {torch.float32: 1e-6, torch.float16: 2e-3}[dtype]
      rtol = {torch.float32: 1e-5, torch.float16: 1e-2}[dtype]
      q_assert_close(bmm, qbmm, atol=atol, rtol=rtol)
  
imports: import pytest
  import torch
  from helpers import q_assert_close, random_qtensor, random_tensor
  from quanto import QTensor
label_desc: Use image node_id 2 for calling a global function with example usage: test_bmm(dtype, batch_size, a_shape, b_shape, b_axis, device) without return types