code: def test_forward_backward(self):
      for (
          dtype,
          seq_length,
          hidden_size,
          rotary_percent,
          transpose,
          transpose_output_memory,
          loss_func,
          cached,
      ) in itertools.product(
          self.dtype,
          self.seq_length,
          self.hidden_size,
          self.rotary_percent,
          self.transpose,
          self.transpose_output_memory,
          self.loss_func,
          self.cached,
      ):
          t = torch.rand(
              (seq_length, self.batch_size, 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(
              (seq_length, 1, 1, int(hidden_size * rotary_percent)),
              dtype=torch.float32,
              device=self.device,
          )
  
          # unfused
          output_unfused = apply_rotary_pos_emb(t, emb)
          loss_unfused = loss_func(output_unfused)
          loss_unfused.backward()
          grad_unfused = t.grad.detach().clone()
          t.grad = None
  
          # fused
          if cached:
              cos, sin = emb.cos(), emb.sin()
              output_fused = fused_apply_rotary_pos_emb_cached(
                  t,
                  cos,
                  sin,
                  transpose_output_memory=transpose_output_memory,
              )
          else:
              output_fused = fused_apply_rotary_pos_emb(
                  t,
                  emb,
                  transpose_output_memory=transpose_output_memory,
              )
          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=}, {seq_length=}, {hidden_size=}, {rotary_percent=}, "
              f"{transpose=}, {transpose_output_memory=}, loss_func={loss_func.__name__}",
          )
          self.assertEqual(
              grad_unfused,
              grad_fused,
              msg=f"{dtype=}, {seq_length=}, {hidden_size=}, {rotary_percent=}, "
              f"{transpose=}, {transpose_output_memory=}, loss_func={loss_func.__name__}",
          )
          assert (
              output_fused.transpose(0, 1).is_contiguous()
              is transpose_output_memory
          )
  
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 5 for calling the TestFusedRoPE obj's underlying member method code with example usage: obj.test_forward_backward() without return types