code: def sdeint_adjoint(
      sde: nn.Module,
      y0: Tensor,
      ts: Vector,
      bm: Optional[BaseBrownian] = None,
      method: Optional[str] = None,
      adjoint_method: Optional[str] = None,
      dt: Scalar = 1e-3,
      adaptive: bool = False,
      adjoint_adaptive: bool = False,
      rtol: Scalar = 1e-5,
      adjoint_rtol: Scalar = 1e-5,
      atol: Scalar = 1e-4,
      adjoint_atol: Scalar = 1e-4,
      dt_min: Scalar = 1e-5,
      options: Optional[Dict[str, Any]] = None,
      adjoint_options: Optional[Dict[str, Any]] = None,
      adjoint_params=None,
      names: Optional[Dict[str, str]] = None,
      logqp: bool = False,
      extra: bool = False,
      extra_solver_state: Optional[Tensors] = None,
      **unused_kwargs,
  ) -> TensorOrTensors:
      """Numerically integrate an SDE with stochastic adjoint support.
  
      Args:
          sde (torch.nn.Module): Object with methods `f` and `g` representing the
              drift and diffusion. The output of `g` should be a single tensor of
              size (batch_size, d) for diagonal noise SDEs or (batch_size, d, m)
              for SDEs of other noise types; d is the dimensionality of state and
              m is the dimensionality of Brownian motion.
          y0 (Tensor): A tensor for the initial state.
          ts (Tensor or sequence of float): Query times in non-descending order.
              The state at the first time of `ts` should be `y0`.
          bm (Brownian, optional): A 'BrownianInterval', `BrownianPath` or
              `BrownianTree` object. Should return tensors of size (batch_size, m)
              for `__call__`. Defaults to `BrownianInterval`.
          method (str, optional): Numerical integration method to use. Must be
              compatible with the SDE type (Ito/Stratonovich) and the noise type
              (scalar/additive/diagonal/general). Defaults to a sensible choice
              depending on the SDE type and noise type of the supplied SDE.
          adjoint_method (str, optional): Name of numerical integration method for
              backward adjoint solve. Defaults to a sensible choice depending on
              the SDE type and noise type of the supplied SDE.
          dt (float, optional): The constant step size or initial step size for
              adaptive time-stepping.
          adaptive (bool, optional): If `True`, use adaptive time-stepping.
          adjoint_adaptive (bool, optional): If `True`, use adaptive time-stepping
              for the backward adjoint solve.
          rtol (float, optional): Relative tolerance.
          adjoint_rtol (float, optional): Relative tolerance for backward adjoint
              solve.
          atol (float, optional): Absolute tolerance.
          adjoint_atol (float, optional): Absolute tolerance for backward adjoint
              solve.
          dt_min (float, optional): Minimum step size during integration.
          options (dict, optional): Dict of options for the integration method.
          adjoint_options (dict, optional): Dict of options for the integration
              method of the backward adjoint solve.
          adjoint_params (Sequence of Tensors, optional): Tensors whose gradient
              should be obtained with the adjoint. If not specified, defaults to
              the parameters of `sde`.
          names (dict, optional): Dict of method names for drift and diffusion.
              Expected keys are "drift" and "diffusion". Serves so that users can
              use methods with names not in `("f", "g")`, e.g. to use the
              method "foo" for the drift, we supply `names={"drift": "foo"}`.
          logqp (bool, optional): If `True`, also return the log-ratio penalty.
          extra (bool, optional): If `True`, also return the extra hidden state
              used internally in the solver.
          extra_solver_state: (tuple of Tensors, optional): Additional state to
              initialise the solver with. Some solvers keep track of additional
              state besides y0, and this offers a way to optionally initialise
              that state.
  
      Returns:
          A single state tensor of size (T, batch_size, d).
          if logqp is True, then the log-ratio penalty is also returned.
          If extra is True, the any extra internal state of the solver is also
          returned.
  
      Raises:
          ValueError: An error occurred due to unrecognized noise type/method,
              or `sde` is missing required methods.
  
      Note:
          The backward pass is much more efficient with Stratonovich SDEs than
          with Ito SDEs.
  
      Note:
          Double-backward is supported for Stratonovich SDEs. Doing so will use
          the adjoint method to compute the gradient of the adjoint. (i.e. rather
          than backpropagating through the numerical solver used for the
          adjoint.) The same `adjoint_method`, `adjoint_adaptive`, `adjoint_rtol,
          `adjoint_atol`, `adjoint_options` will be used for the second-order
          adjoint as is used for the first-order adjoint.
      """
      misc.handle_unused_kwargs(unused_kwargs, msg="`sdeint_adjoint`")
      del unused_kwargs
  
      if adjoint_params is None and not isinstance(sde, nn.Module):
          raise ValueError(
              "`sde` must be an instance of nn.Module to specify the adjoint parameters; alternatively they "
              "can be specified explicitly via the `adjoint_params` argument. If there are no parameters "
              "then it is allowable to set `adjoint_params=()`."
          )
  
      sde, y0, ts, bm, method, options = sdeint.check_contract(
          sde, y0, ts, bm, method, adaptive, options, names, logqp
      )
      misc.assert_no_grad(
          [
              "ts",
              "dt",
              "rtol",
              "adjoint_rtol",
              "atol",
              "adjoint_atol",
              "dt_min",
          ],
          [ts, dt, rtol, adjoint_rtol, atol, adjoint_atol, dt_min],
      )
      adjoint_params = (
          tuple(sde.parameters())
          if adjoint_params is None
          else tuple(adjoint_params)
      )
      adjoint_params = filter(lambda x: x.requires_grad, adjoint_params)
      adjoint_method = _select_default_adjoint_method(
          sde, method, adjoint_method
      )
      adjoint_options = (
          {} if adjoint_options is None else adjoint_options.copy()
      )
  
      # Note that all of these warnings are only applicable for reversible solvers with sdeint_adjoint; none of them
      # apply to sdeint.
      if method == METHODS.reversible_heun:
          if adjoint_method != METHODS.adjoint_reversible_heun:
              warnings.warn(
                  f"method={repr(method)}, but adjoint_method!={repr(METHODS.adjoint_reversible_heun)}."
              )
          if adaptive or adjoint_adaptive:
              warnings.warn(
                  f"A limitation of the current method={repr(method)} implementation is "
                  f"that it does not save the time steps used. This means that it may not be perfectly "
                  f"accurate when used with `adaptive` or `adjoint_adaptive`."
              )
          else:
              num_steps = (ts - ts[0]) / dt
              if not torch.allclose(num_steps, num_steps.round()):
                  warnings.warn(
                      f"The spacing between time points `ts` is not an integer multiple of the time step `dt`. "
                      f"This means that the backward pass (which is forced to step to each of `ts` to get "
                      f"dL/dy(t) for t in ts) will not perfectly mimick the forward pass (which does not step "
                      f"to each `ts`, and instead interpolates to them). This means that "
                      f"method={repr(method)} may not be perfectly accurate."
                  )
  
      solver_fn = methods.select(method=method, sde_type=sde.sde_type)
      solver = solver_fn(
          sde=sde,
          bm=bm,
          dt=dt,
          adaptive=adaptive,
          rtol=rtol,
          atol=atol,
          dt_min=dt_min,
          options=options,
      )
      if extra_solver_state is None:
          extra_solver_state = solver.init_extra_solver_state(ts[0], y0)
  
      ys, *extra_solver_state = _SdeintAdjointMethod.apply(
          sde,
          ts,
          dt,
          bm,
          solver,
          method,
          adjoint_method,
          adjoint_adaptive,
          adjoint_rtol,
          adjoint_atol,
          dt_min,
          adjoint_options,
          len(extra_solver_state),
          y0,
          *extra_solver_state,
          *adjoint_params,
      )
  
      return sdeint.parse_return(
          y0, ys, extra_solver_state, extra, logqp
      )
  
file_path: torchsde/torchsde/_core/adjoint.py
imports: import torch
  from torch import nn
  import warnings
  from .None import base_sde
  from .None import methods
  from .None import misc
  from .None import sdeint
  from .adjoint_sde import AdjointSDE
  from .._brownian import BaseBrownian, ReverseBrownian
  from ..settings import METHODS, NOISE_TYPES, SDE_TYPES
  from ..types import Any, Dict, Optional, Scalar, Tensor, Tensors, TensorOrTensors, Vector
label_desc: Use image node_id 1 for calling a global function with example usage: sdeint_adjoint(sde, y0, ts, bm, method, adjoint_method, dt, adaptive, adjoint_adaptive, rtol, adjoint_rtol, atol, adjoint_atol, dt_min, options, adjoint_options, adjoint_params, names, logqp, extra, extra_solver_state) and returns: sdeint