code: def step(self, t0, t1, y0, extra0):
      forward_f0, forward_g0, forward_z0 = extra0
      dt = t1 - t0
      dW = self.bm(t0, t1)
      half_dt = 0.5 * dt
      half_dW = 0.5 * dW
      (
          forward_y0,
          adj_y0,
          (adj_f0, adj_g0, adj_z0, *adj_params),
          requires_grad,
      ) = self.sde.get_state(t0, y0, extra_states=True)
      adj_y0_half_dt = adj_y0 * half_dt
      adj_y0_half_dW = self._adjoint_of_prod(adj_y0, half_dW)
  
      forward_z1 = (
          2 * forward_y0
          - forward_z0
          - forward_f0 * dt
          - self.forward_sde.prod(forward_g0, dW)
      )
  
      adj_y1 = adj_y0
      adj_f1 = adj_y0_half_dt
      adj_f0 = adj_f0 + adj_y0_half_dt
      adj_g1 = adj_y0_half_dW
      adj_g0 = adj_g0 + adj_y0_half_dW
  
      # TODO: efficiency. It should be possible to make one fewer forward call by re-using the forward computation
      #  in the previous step.
      with torch.enable_grad():
          if not forward_z0.requires_grad:
              forward_z0 = forward_z0.detach().requires_grad_()
          re_forward_f0, re_forward_g0 = self.forward_sde.f_and_g(
              -t0, forward_z0
          )
  
          vjp_z, *vjp_params = misc.vjp(
              outputs=(re_forward_f0, re_forward_g0),
              inputs=[forward_z0] + self.sde.params,
              grad_outputs=[adj_f0, adj_g0],
              allow_unused=True,
              retain_graph=True,
              create_graph=requires_grad,
          )
      adj_z0 = adj_z0 + vjp_z
      adj_params = misc.seq_add(adj_params, vjp_params)
  
      forward_f1, forward_g1 = self.forward_sde.f_and_g(-t1, forward_z1)
      forward_y1 = (
          forward_y0
          - (forward_f0 + forward_f1) * half_dt
          - self.forward_sde.prod(forward_g0 + forward_g1, half_dW)
      )
  
      adj_y1 = adj_y1 + 2 * adj_z0
      adj_z1 = -adj_z0
      adj_f1 = adj_f1 + adj_z0 * dt
      adj_g1 = adj_g1 + self._adjoint_of_prod(adj_z0, dW)
  
      y1 = misc.flatten(
          [forward_y1, adj_y1, adj_f1, adj_g1, adj_z1] + adj_params
      ).unsqueeze(0)
  
      return y1, (forward_f1, forward_g1, forward_z1)
  
file_path: torchsde/torchsde/_core/methods/reversible_heun.py
imports: import torch
  from ..None import adjoint_sde
  from ..None import base_solver
  from ..None import misc
  from ...settings import SDE_TYPES, NOISE_TYPES, LEVY_AREA_APPROXIMATIONS, METHODS
label_desc: Use image node_id 3 for calling the AdjointReversibleHeun obj's underlying member method code with example usage: obj.step(t0, t1, y0, extra0) and returns: y1,  forward_f1,  forward_g1,  forward_z1