code: def test_adjoint_norm(self):
      def f(t, x):
          return x
  
      t = torch.tensor([0.0, 1.0])
      adjoint_params = (
          torch.rand(7, requires_grad=True),
          torch.rand((), requires_grad=True),
      )
  
      def make_spy_on_adjoint_norm(adjoint_norm, actual_norm):
          is_spy_called = [False]
  
          def spy_on_adjoint_norm(tensor):
              nonlocal is_spy_called
              is_spy_called[0] = True
              norm_result = adjoint_norm(tensor)
              true_norm_result = actual_norm(tensor)
              self.assertIsInstance(norm_result, torch.Tensor)
              self.assertEqual(
                  norm_result.shape, true_norm_result.shape
              )
              self.assertLess(
                  (norm_result - true_norm_result).abs().max(), 1e-6
              )
              return norm_result
  
          return spy_on_adjoint_norm, is_spy_called
  
      # Test the various auto-constructed adjoint norms with tensor (not tuple) state
      for shape in ((), (1,), (2, 2)):
          for use_adjoint_options, seminorm in (
              (False, False),
              (True, False),
              (True, True),
          ):
              with self.subTest(
                  shape=shape,
                  use_adjoint_options=use_adjoint_options,
                  seminorm=seminorm,
              ):
                  x0 = torch.full(shape, 1.0)
                  if use_adjoint_options:
                      if seminorm:
                          # Test passing adjoint_options and wanting the seminorm
                          kwargs = dict(
                              adjoint_options=dict(norm="seminorm")
                          )
                      else:
                          # Test passing adjoint_options but not specify the adjoint norm
                          kwargs = dict(adjoint_options={})
                  else:
                      # Test not passing adjoint_options at all.
                      kwargs = {}
                  xs = torchdiffeq.odeint_adjoint(
                      f, x0, t, adjoint_params=adjoint_params, **kwargs
                  )
                  _adjoint_norm = xs.grad_fn.adjoint_options["norm"]
  
                  is_called = False
  
                  def actual_norm(tensor_tuple):
                      nonlocal is_called
                      is_called = True
                      self.assertIsInstance(tensor_tuple, tuple)
                      t, y, adj_y, adj_param1, adj_param2 = tensor_tuple
                      self.assertEqual(t.shape, ())
                      self.assertEqual(y.shape, shape)
                      self.assertEqual(adj_y.shape, shape)
                      self.assertEqual(adj_param1.shape, (7,))
                      self.assertEqual(adj_param2.shape, (()))
                      out = max(
                          t.abs(),
                          y.pow(2).mean().sqrt(),
                          adj_y.pow(2).mean().sqrt(),
                      )
                      if not seminorm:
                          out = max(
                              out,
                              adj_param1.pow(2).mean().sqrt(),
                              adj_param2.abs(),
                          )
                      return out
  
                  (
                      xs.grad_fn.adjoint_options["norm"],
                      is_spy_called,
                  ) = make_spy_on_adjoint_norm(
                      _adjoint_norm, actual_norm
                  )
                  xs.sum().backward()
                  self.assertTrue(is_called)
                  self.assertTrue(is_spy_called[0])
  
      # Test the various auto-constructed adjoint norms with tuple (not tensor) state
      for use_adjoint_options, seminorm in (
          (False, False),
          (True, False),
          (True, True),
      ):
          with self.subTest(
              shape=shape,
              use_adjoint_options=use_adjoint_options,
              seminorm=seminorm,
          ):
              x0 = torch.tensor(1.0), torch.tensor(
                  [[0.5, 0.5], [0.1, 0.1]]
              )
              if use_adjoint_options:
                  if seminorm:
                      # Test passing adjoint_options and wanting the seminorm
                      kwargs = dict(
                          adjoint_options=dict(norm="seminorm")
                      )
                  else:
                      # Test passing adjoint_options but not specify the adjoint norm
                      kwargs = dict(adjoint_options={})
              else:
                  # Test not passing adjoint_options at all.
                  kwargs = {}
              xs = torchdiffeq.odeint_adjoint(
                  f, x0, t, adjoint_params=adjoint_params, **kwargs
              )
              adjoint_options_dict = (
                  xs[0]
                  .grad_fn.next_functions[0][0]
                  .next_functions[0][0]
                  .adjoint_options
              )
              _adjoint_norm = adjoint_options_dict["norm"]
  
              is_called = False
  
              def actual_norm(tensor_tuple):
                  nonlocal is_called
                  is_called = True
                  self.assertIsInstance(tensor_tuple, tuple)
                  t, y, adj_y, adj_param1, adj_param2 = tensor_tuple
                  self.assertEqual(t.shape, ())
                  self.assertEqual(y.shape, (5,))
                  self.assertEqual(adj_y.shape, (5,))
                  self.assertEqual(adj_param1.shape, (7,))
                  self.assertEqual(adj_param2.shape, ())
                  ya = y[0]
                  yb = y[1:]
                  adj_ya = adj_y[0]
                  adj_yb = adj_y[1:4]
                  out = max(
                      t.abs(),
                      ya.abs(),
                      yb.pow(2).mean().sqrt(),
                      adj_ya.abs(),
                      adj_yb.pow(2).mean().sqrt(),
                  )
                  if not seminorm:
                      out = max(
                          out,
                          adj_param1.pow(2).mean().sqrt(),
                          adj_param2.abs(),
                      )
                  return out
  
              (
                  spy_on_adjoint_norm,
                  is_spy_called,
              ) = make_spy_on_adjoint_norm(_adjoint_norm, actual_norm)
              adjoint_options_dict["norm"] = spy_on_adjoint_norm
              xs[0].sum().backward()
              self.assertTrue(is_called)
              self.assertTrue(is_spy_called[0])
  
      # Test user-passed adjoint norms with tensor (not tuple) state
      is_called = False
  
      def adjoint_norm(tensor_tuple):
          nonlocal is_called
          is_called = True
          self.assertIsInstance(tensor_tuple, tuple)
          t, y, adj_y, adj_param1, adj_param2 = tensor_tuple
          self.assertEqual(t.shape, ())
          self.assertEqual(y.shape, ())
          self.assertEqual(adj_y.shape, ())
          self.assertEqual(adj_param1.shape, (7,))
          self.assertEqual(adj_param2.shape, ())
          return max(
              t.abs(),
              y.pow(2).mean().sqrt(),
              adj_y.pow(2).mean().sqrt(),
              adj_param1.pow(2).mean().sqrt(),
              adj_param2.abs(),
          )
  
      x0 = torch.tensor(1.0)
      xs = torchdiffeq.odeint_adjoint(
          f,
          x0,
          t,
          adjoint_params=adjoint_params,
          adjoint_options=dict(norm=adjoint_norm),
      )
      xs.sum().backward()
      self.assertTrue(is_called)
  
      # Test user-passed adjoint norms with tuple (not tensor) state
      is_called = False
  
      def adjoint_norm(tensor_tuple):
          nonlocal is_called
          is_called = True
          self.assertIsInstance(tensor_tuple, tuple)
          (
              t,
              ya,
              yb,
              adj_ya,
              adj_yb,
              adj_param1,
              adj_param2,
          ) = tensor_tuple
          self.assertEqual(t.shape, ())
          self.assertEqual(ya.shape, ())
          self.assertEqual(yb.shape, (2, 2))
          self.assertEqual(adj_ya.shape, ())
          self.assertEqual(adj_yb.shape, (2, 2))
          self.assertEqual(adj_param1.shape, (7,))
          self.assertEqual(adj_param2.shape, ())
          return max(
              t.abs(),
              ya.abs(),
              yb.pow(2).mean().sqrt(),
              adj_ya.abs(),
              adj_yb.pow(2).mean().sqrt(),
              adj_param1.pow(2).mean().sqrt(),
              adj_param2.abs(),
          )
  
      x0 = torch.tensor(1.0), torch.tensor([[0.5, 0.5], [0.1, 0.1]])
      xs = torchdiffeq.odeint_adjoint(
          f,
          x0,
          t,
          adjoint_params=adjoint_params,
          adjoint_options=dict(norm=adjoint_norm),
      )
      xs[0].sum().backward()
      self.assertTrue(is_called)
  
file_path: torchdiffeq/tests/norm_tests.py
imports: import contextlib
  import unittest
  import torch
  import torchdiffeq
  from problems import DTYPES, DEVICES, ADAPTIVE_METHODS
label_desc: Use image node_id 2 for calling the TestNorms obj's underlying member method code with example usage: obj.test_adjoint_norm() and returns: x, spy_on_adjoint_norm,  is_spy_called, max, max, norm_result, out, out