dew.sampling.flow

Stochastic rectified-flow transitions and their Gaussian likelihoods.

Flow-GRPO (arXiv:2505.05470v5, equations 8-9), read against https://github.com/yifan123/flow_grpo/blob/879042cf5707f8b90daa98d147d7deac2317c5da/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py uses diffusion coefficient a * sqrt(t / (1 - t)). At t=1, the reference replaces the denominator’s time with the first interior grid point.

Name	Summary
FlowSDE	Flow-GRPO’s Euler-Maruyama solver on a rectified-flow Process.
FlowTrajectory	A reverse trajectory, with batch-major states and joint log densities.
GaussianTransition	An isotropic transition with one variance per batch row.
flow_transition	Euler-Maruyama over the physical noise rate, with 0 <= sigma_next <= sigma <= 1.
sample_trajectory	Record FlowSDE transitions over the same time grid and keys as sample.

FlowSDE

dataclass source

class FlowSDE(noise_level: float = 0.7)

Flow-GRPO’s Euler-Maruyama solver on a rectified-flow Process.

Process times may be resolution-shifted. The transition integrates in the resulting physical noise rate, as the reference scheduler does.

FlowSDE.validate

def validate(process: Process) -> None

FlowSDE.init

def init(
    x: jax.Array,
    times: jax.Array,
    process: Process,
    *,
    key: jax.Array,
) -> tuple[()]

FlowSDE.transition

def transition(
    x: jax.Array,
    t: jax.Array,
    t_next: jax.Array,
    denoised: jax.Array,
    eps: jax.Array,
    process: Process,
) -> GaussianTransition

FlowSDE.step

def step(
    x: jax.Array,
    t: jax.Array,
    t_next: jax.Array,
    denoised: jax.Array,
    eps: jax.Array,
    state: tuple[()],
    key: jax.Array,
    process: Process,
    denoise: Callable[[jax.Array, jax.Array], tuple[jax.Array, jax.Array]],
    /,
) -> tuple[jax.Array, tuple[()]]

FlowTrajectory

dataclass source

class FlowTrajectory()

A reverse trajectory, with batch-major states and joint log densities.

states is [batch, points, …], times is [points], and log_probs and stochastic are [batch, points - 1]. Deterministic intervals have NaN log density and a false stochastic mark. The final state is the sample.

GaussianTransition

dataclass source

class GaussianTransition()

An isotropic transition with one variance per batch row.

Sampling and density arithmetic use float32. Densities and KL sum over the sample dimensions. A zero variance is a deterministic transition: sampling returns its mean and log_prob is NaN, since a Dirac measure has no density with respect to Lebesgue measure.

GaussianTransition.sample

def sample(key: jax.Array) -> jax.Array

GaussianTransition.log_prob

def log_prob(value: ArrayLike) -> jax.Array

Joint log density of an observed next state, one value per row.

GaussianTransition.kl

def kl(reference_mean: ArrayLike) -> jax.Array

KL to a reference transition with the same policy-independent variance.

Equal Dirac measures have KL zero; distinct ones have infinite KL.

flow_transition

function source

def flow_transition(
    x: ArrayLike,
    velocity: ArrayLike,
    sigma: ArrayLike,
    sigma_next: ArrayLike,
    *,
    noise_level: float = 0.7,
) -> GaussianTransition

Euler-Maruyama over the physical noise rate, with 0 <= sigma_next <= sigma <= 1.

A rectified flow’s noise rate is its own physical time, which is what FlowSDE reads off the schedule and hands over here. x and velocity are [batch, …]; rates are scalars or [batch]. All density arithmetic is float32. Variance is sigma^2 times the elapsed rate. Invalid rates produce non-finite transitions. At zero noise or zero elapsed rate the result is deterministic.

sample_trajectory

function source

def sample_trajectory(
    denoise: Denoiser,
    x_T: jax.Array,
    steps: int,
    *,
    solver: FlowSDE = FlowSDE(),
    guidance: CFG | None = None,
    key: jax.Array,
) -> FlowTrajectory

Record FlowSDE transitions over the same time grid and keys as sample.

steps counts grid points, including both endpoints. A ten-transition rollout therefore uses steps=11. Guidance is applied identically before constructing each Gaussian. Rectified flow’s clean prediction at t=0 is its state, so the last transition already produces the final sample.