code: def test_DS4Sci_EvoformerAttention(dtype, tensor_shape):
      skip_on_arch(8 if dtype == torch.bfloat16 else 7)
      batch, n, seq_len, heads, dim = tensor_shape
      Q = torch.randn(
          batch,
          n,
          seq_len,
          heads,
          dim,
          dtype=dtype,
          device=get_accelerator().device_name(),
          requires_grad=True,
      )
      K = torch.randn(
          batch,
          n,
          seq_len,
          heads,
          dim,
          dtype=dtype,
          device=get_accelerator().device_name(),
          requires_grad=True,
      )
      V = torch.randn(
          batch,
          n,
          seq_len,
          heads,
          dim,
          dtype=dtype,
          device=get_accelerator().device_name(),
          requires_grad=True,
      )
      mask = torch.randint(
          0,
          2,
          (batch, n, 1, 1, seq_len),
          dtype=dtype,
          device=get_accelerator().device_name(),
      )
      mask_bias = 1e9 * (mask - 1)
      bias = torch.randn(
          batch,
          1,
          heads,
          seq_len,
          seq_len,
          dtype=dtype,
          device=get_accelerator().device_name(),
          requires_grad=True,
      )
      dummy_out = torch.rand_like(
          Q, dtype=dtype, device=get_accelerator().device_name()
      )
      ref_out = attention_reference(
          Q, K, V, [mask_bias, bias], 1 / (dim**0.5)
      )
      ref_out.backward(dummy_out)
      ref_dv, V.grad = V.grad.clone(), None
      ref_dk, K.grad = K.grad.clone(), None
      ref_dq, Q.grad = Q.grad.clone(), None
      ref_db, bias.grad = bias.grad.clone(), None
  
      out = DS4Sci_EvoformerAttention(Q, K, V, [mask_bias, bias])
      out.backward(dummy_out)
      dv, v_grad = V.grad.clone(), None
      dk, k_grad = K.grad.clone(), None
      dq, q_grad = Q.grad.clone(), None
      db, bias.grad = bias.grad.clone(), None
  
      eps = 1e-2 if dtype == torch.float16 else 5e-2
  
      assert (
          torch.max(torch.abs(ref_out - out)).item() < eps
      ), f"out eps: {torch.max(torch.abs(ref_out - out))}"
      assert (
          torch.max(torch.abs(ref_dv - dv)) < eps
      ), f"dv eps: {torch.max(torch.abs(ref_dv - dv))}"
      assert (
          torch.max(torch.abs(ref_dk - dk)) < eps
      ), f"dk eps: {torch.max(torch.abs(ref_dk - dk))}"
      assert (
          torch.max(torch.abs(ref_dq - dq)) < eps
      ), f"dq eps: {torch.max(torch.abs(ref_dq - dq))}"
      assert (
          torch.max(torch.abs(ref_db - db)) < 2 * eps
      ), f"db eps: {torch.max(torch.abs(ref_db - db))}"
  
file_path: DeepSpeed/tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py
imports: from typing import List
  import pytest
  import torch
  from torch.nn import functional
  import deepspeed
  from deepspeed.ops.op_builder import EvoformerAttnBuilder
  from deepspeed.ops.deepspeed4science import DS4Sci_EvoformerAttention
  from deepspeed.accelerator import get_accelerator
  from unit.util import skip_on_arch
label_desc: Use image node_id 2 for calling a global function with example usage: test_DS4Sci_EvoformerAttention(dtype, tensor_shape) without return types