code: def test_multiple_subgraphs(self, dtype, device):
      # Checks that rng state is maintained when there are multiple aot traced
      # graphs.
      shape = (16, 16)
  
      class CustomOp1(torch.autograd.Function):
          @staticmethod
          def forward(ctx, x):
              ctx.save_for_backward(x)
              a = torch.rand_like(x) * x
              a = torch.rand_like(x) * a
              return a
  
          @staticmethod
          def backward(ctx, grad_out):
              (x,) = ctx.saved_tensors
              return grad_out * torch.rand_like(grad_out) * torch.cos(x)
  
      class CustomOp2(torch.autograd.Function):
          @staticmethod
          def forward(ctx, x):
              ctx.save_for_backward(x)
              a = torch.rand_like(x) * x
              return a
  
          @staticmethod
          def backward(ctx, grad_out):
              (x,) = ctx.saved_tensors
              return (
                  grad_out
                  * torch.rand_like(grad_out)
                  * torch.rand_like(x)
              )
  
      custom_op1 = CustomOp1.apply
      custom_op2 = CustomOp2.apply
  
      def fn(x):
          a = custom_op1(x)
          b = a.sin()
          return custom_op2(b)
  
      fwd_compiler = functools.partial(count_philox_rand, freq=2)
      bwd_compiler = functools.partial(count_philox_rand, freq=1)
      aot_custom_op1 = aot_function(
          custom_op1, fwd_compiler, bwd_compiler
      )
      fwd_compiler = functools.partial(count_philox_rand, freq=1)
      bwd_compiler = functools.partial(count_philox_rand, freq=2)
      aot_custom_op2 = aot_function(
          custom_op2, fwd_compiler, bwd_compiler
      )
  
      def aot_fn(x):
          a = aot_custom_op1(x)
          b = a.sin()
          return aot_custom_op2(b)
  
      for seed in range(10):
          torch.cuda.manual_seed(seed)
          x = torch.rand(
              *shape, device=device, dtype=dtype, requires_grad=True
          )
          x_clone = x.clone().detach().requires_grad_(True)
  
          torch.cuda.manual_seed(seed)
          ref = fn(x)
          ref.sum().backward()
  
          torch.cuda.manual_seed(seed)
          res = aot_fn(x_clone)
          res.sum().backward()
  
          self.assertEqual(ref, res)
          self.assertEqual(x.grad, x_clone.grad)
  
imports: import sys
  import unittest
  import torch
  from torch.testing._internal.common_utils import TestCase, run_tests
  from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes
  from functorch.compile import aot_function, nop, min_cut_rematerialization_partition
  from unittest.mock import patch
  import functools
  import torch.utils.checkpoint
  from torch.testing._internal.common_utils import IS_CI, IS_WINDOWS
label_desc: Use image node_id 6 for calling the TestFunctionalizationRngOps obj's underlying member method code with example usage: obj.test_multiple_subgraphs(dtype, device) and returns: custom_op2, aot_custom_op2, a, unknown, a, unknown