def reverse()

in jcm/sde_lib.py [0:0]

• 28 lines of code
• 7 McCabe index (conditional complexity)

    def reverse(self, score_fn, probability_flow=False):
        """Create the reverse-time SDE/ODE.

        Args:
          score_fn: A time-dependent score-based model that takes x and t and returns the score.
          probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
        """
        N = self.N
        T = self.T
        sde_fn = self.sde
        discretize_fn = self.discretize

        # Build the class for reverse-time SDE.
        class RSDE(self.__class__):
            def __init__(self):
                self.N = N
                self.probability_flow = probability_flow

            @property
            def T(self):
                return T

            def sde(self, x, t):
                """Create the drift and diffusion functions for the reverse SDE/ODE."""
                drift, diffusion = sde_fn(x, t)
                score = score_fn(x, t)
                drift = drift - batch_mul(
                    diffusion**2, score * (0.5 if self.probability_flow else 1.0)
                )
                # Set the diffusion function to zero for ODEs.
                diffusion = jnp.zeros_like(t) if self.probability_flow else diffusion
                return drift, diffusion

            def discretize(self, x, t):
                """Create discretized iteration rules for the reverse diffusion sampler."""
                f, G = discretize_fn(x, t)
                rev_f = f - batch_mul(
                    G**2, score_fn(x, t) * (0.5 if self.probability_flow else 1.0)
                )
                rev_G = jnp.zeros_like(t) if self.probability_flow else G
                return rev_f, rev_G

        return RSDE()