docstr: Unpacks y_aug, whilst enforcing the necessary checks so that we can calculate derivatives wrt state.
docstr_tok: ["Unpacks","y_aug",",","whilst","enforcing","the","necessary","checks","so","that","we","can","calculate","derivatives","wrt","state","."]
code: def get_state(self, t, y_aug, v=None, extra_states=False):
      """Unpacks y_aug, whilst enforcing the necessary checks so that we can calculate derivatives wrt state."""
  
      # These leaf checks are very important.
      # get_state is used where we want to compute:
      # ```
      # with torch.enable_grad():
      #     s = some_function(y)
      #     torch.autograd.grad(s, [y] + params, ...)
      # ```
      # where `some_function` implicitly depends on `params`.
      # However if y has history of its own then in principle it could _also_ depend upon `params`, and this call to
      # `grad` will go all the way back to that. To avoid this, we require that every input tensor be a leaf tensor.
      #
      # This is also the reason for the `y0.detach()` in adjoint.py::_SdeintAdjointMethod.forward. If we don't detach,
      # then y0 may have a history and these checks will fail. This is a spurious failure as
      # `torch.autograd.Function.forward` has an implicit `torch.no_grad()` guard, i.e. we definitely don't want to
      # use its history there.
      assert (
          t.is_leaf
      ), "Internal error: please report a bug to torchsde"
      assert (
          y_aug.is_leaf
      ), "Internal error: please report a bug to torchsde"
      if v is not None:
          assert (
              v.is_leaf
          ), "Internal error: please report a bug to torchsde"
  
      requires_grad = torch.is_grad_enabled()
  
      if extra_states:
          shapes = self._shapes
      else:
          shapes = self._shapes[:2]
      numel = sum(shape.numel() for shape in shapes)
      y, adj_y, *extra_states = misc.flat_to_shape(
          y_aug.squeeze(0)[:numel], shapes
      )
  
      # To support the later differentiation wrt y, we set it to require_grad if it doesn't already.
      if not y.requires_grad:
          y = y.detach().requires_grad_()
      return y, adj_y, extra_states, requires_grad
  
imports: import torch
  from .None import base_sde
  from .None import misc
  from ..settings import NOISE_TYPES, SDE_TYPES
  from ..types import Sequence, TensorOrTensors
label_desc: Use image node_id 2 for calling the AdjointSDE obj's underlying member method code with example usage: obj.get_state(t, y_aug, v, extra_states) and returns: y,  adj_y,  extra_states,  requires_grad