class_docstr: SDE abstract class. Functions are designed for a mini-batch of inputs.
docstr: Create the reverse-time SDE/ODE.
  
  Args:
      score_model: A function that takes x, t and y and returns the score.
      probability_flow: If `True`, create the reverse-time ODE
          used for probability flow sampling.
code: def reverse(oself, score_model, probability_flow=False):
      """Create the reverse-time SDE/ODE.
  
      Args:
          score_model: A function that takes x, t and y and returns the score.
          probability_flow: If `True`, create the reverse-time ODE
              used for probability flow sampling.
      """
      N = oself.N
      T = oself.T
      sde_fn = oself.sde
      discretize_fn = oself.discretize
  
      # Build the class for reverse-time SDE.
      class RSDE(oself.__class__):
          def __init__(self):
              self.N = N
              self.probability_flow = probability_flow
  
          @property
          def T(self):
              return T
  
          def sde(self, x, t, *args):
              """Create the drift and diffusion functions for the reverse SDE/ODE."""
              rsde_parts = self.rsde_parts(x, t, *args)
              total_drift, diffusion = (
                  rsde_parts["total_drift"],
                  rsde_parts["diffusion"],
              )
              return total_drift, diffusion
  
          def rsde_parts(self, x, t, *args):
              sde_drift, sde_diffusion = sde_fn(x, t, *args)
              score = score_model(x, t, *args)
              score_drift = (
                  -sde_diffusion[:, None, None, None] ** 2
                  * score
                  * (0.5 if self.probability_flow else 1.0)
              )
              diffusion = (
                  torch.zeros_like(sde_diffusion)
                  if self.probability_flow
                  else sde_diffusion
              )
              total_drift = sde_drift + score_drift
              return {
                  "total_drift": total_drift,
                  "diffusion": diffusion,
                  "sde_drift": sde_drift,
                  "sde_diffusion": sde_diffusion,
                  "score_drift": score_drift,
                  "score": score,
              }
  
          def discretize(self, x, t, *args):
              """Create discretized iteration rules for the reverse
              diffusion sampler."""
              f, G = discretize_fn(x, t, *args)
              rev_f = f - G[:, None, None, None] ** 2 * score_model(
                  x, t, *args
              ) * (0.5 if self.probability_flow else 1.0)
              rev_G = (
                  torch.zeros_like(G) if self.probability_flow else G
              )
              return rev_f, rev_G
  
      return RSDE()
  
imports: import abc
  import warnings
  import numpy
  import torch
label_desc: Use image node_id 8 for calling the SDE obj's underlying member method code with example usage: obj.reverse(oscore_model, probability_flow) and returns: RSDE, T, total_drift,  diffusion, dict, rev_f,  rev_G