code: def test_thd_forward_backward(self):
      cu_seqlens = torch.tensor(
          [
              0,
              400,
              542,
              711,
              727,
              752,
              1270,
              1426,
              1450,
              1954,
              2044,
              2048,
          ],
          dtype=torch.int32,
          device=self.device,
      )
      for (
          dtype,
          hidden_size,
          rotary_percent,
          transpose,
          loss_func,
      ) in itertools.product(
          self.dtype,
          self.hidden_size,
          self.rotary_percent,
          [None, [1, 2]],
          self.loss_func,
      ):
          t = torch.rand(
              (cu_seqlens[-1], self.head_num, hidden_size),
              dtype=dtype,
              device=self.device,
          )
          if transpose:
              t = (
                  t.transpose(*transpose)
                  .contiguous()
                  .transpose(*transpose)
              )
          t.requires_grad = True
  
          emb = torch.rand(
              (cu_seqlens[-1], 1, 1, int(hidden_size * rotary_percent)),
              dtype=torch.float32,
              device=self.device,
          )
  
          # unfused
          output_unfused = apply_rotary_pos_emb_thd(t, cu_seqlens, emb)
          loss_unfused = loss_func(output_unfused)
          loss_unfused.backward()
          grad_unfused = t.grad.detach().clone()
          t.grad = None
  
          # fused
          output_fused = fused_apply_rotary_pos_emb_thd(
              t,
              cu_seqlens,
              emb,
          )
          loss_fused = loss_func(output_fused)
          loss_fused.backward()
          grad_fused = t.grad.detach().clone()
          t.grad = None
  
          self.assertEqual(
              output_unfused,
              output_fused,
              msg=f"{dtype=}, {cu_seqlens=}, {hidden_size=}, {rotary_percent=}, "
              f"{transpose=}, loss_func={loss_func.__name__}",
          )
          self.assertEqual(
              grad_unfused,
              grad_fused,
              msg=f"{dtype=}, {cu_seqlens=}, {hidden_size=}, {rotary_percent=}, "
              f"{transpose=}, loss_func={loss_func.__name__}",
          )
  
file_path: apex/tests/L0/run_transformer/test_fused_rope.py
imports: import itertools
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.functional import fused_apply_rotary_pos_emb, fused_apply_rotary_pos_emb_cached, fused_apply_rotary_pos_emb_thd
label_desc: Use image node_id 6 for calling the TestFusedRoPE obj's underlying member method code with example usage: obj.test_thd_forward_backward() without return types