code: def test_odeint_jump_t(self):
      for adjoint in (False, True):
          for dtype in DTYPES:
              for device in DEVICES:
                  for method in ADAPTIVE_METHODS:
                      with self.subTest(
                          adjoint=adjoint,
                          dtype=dtype,
                          device=device,
                          method=method,
                      ):
                          if method == "dopri8":
                              # Doesn't seem to work for some reason.
                              continue
  
                          x0 = torch.tensor(
                              [1.0, 2.0],
                              device=device,
                              dtype=dtype,
                              requires_grad=True,
                          )
                          t = torch.tensor([0.0, 1.0], device=device)
  
                          simple_f = _JumpF()
                          odeint = (
                              partial(
                                  torchdiffeq.odeint_adjoint,
                                  adjoint_params=(),
                              )
                              if adjoint
                              else torchdiffeq.odeint
                          )
                          simple_xs = odeint(
                              simple_f, x0, t, atol=1e-6, method=method
                          )
  
                          better_f = _JumpF()
                          options = dict(
                              jump_t=torch.tensor([0.5], device=device)
                          )
  
                          with warnings.catch_warnings():
                              better_xs = odeint(
                                  better_f,
                                  x0,
                                  t,
                                  rtol=1e-6,
                                  atol=1e-6,
                                  method=method,
                                  options=options,
                              )
  
                          self.assertLess(better_f.nfe, simple_f.nfe)
  
                          if adjoint:
                              simple_f.nfe = 0
                              better_f.nfe = 0
                              with warnings.catch_warnings():
                                  simple_xs.sum().backward()
                                  better_xs.sum().backward()
                              self.assertLess(
                                  better_f.nfe, simple_f.nfe
                              )
  
imports: from functools import partial
  import unittest
  import warnings
  import torch
  import torchdiffeq
  from problems import construct_problem, PROBLEMS, DTYPES, DEVICES, METHODS, ADAPTIVE_METHODS, FIXED_METHODS, SCIPY_METHODS
label_desc: Use image node_id 1 for calling the TestDiscontinuities obj's underlying member method code with example usage: obj.test_odeint_jump_t() without return types