Upload 87 files

0ba9d09 verified about 1 month ago

12.4 kB

"""Simulation-Free Score and Flow Matching: two closed-form targets at one bridge sample.
With the entropic plan computed on a minibatch, the bridge is a mixture of Brownian bridges between
pairs drawn from that plan, and a Brownian bridge at a time t is Gaussian. One normal draw gives a
sample, and at that sample both the bridge drift and the score of the bridge marginal are available
without simulating a path. We subtract half the noise variance times the score. That turns the
drift into the velocity of the probability flow, and its dependence on epsilon cancels.
Chapter 7, Section 7.19.3 of the notes. Runs on the endpoint pair x_0 = 0, x_1 = 2 at epsilon = 1,
and on the two-point batch of Section 7.8.1. Prints the two regression targets at t = 0.25 and at
t = 0.75, the scaled score at every time on the grid, and the entropic plan of the batch.
    python lecture_7/sf2m.py              # 101 times on the grid
    python lecture_7/sf2m.py --steps 20   # a coarse time grid, for a first look
"""

# %% 01. Imports and configuration
from __future__ import annotations

import torch

from cis6270.runner import common_parser, print_report, save_run, set_seed

DTYPE = torch.float64  # The chapter quotes four figures of both regression targets.

X0, X1 = 0.0, 2.0  # The endpoint pair the chapter works through.
EPSILON = 1.0  # The reference noise variance per unit time.
DRAW = 0.4  # The normal draw z the chapter uses at both times.


# %% 02. One sample from the bridge
#
# The Brownian bridge at time t is Gaussian with mean the straight interpolation and variance
# epsilon t (1 - t), so a sample costs one normal draw and no intermediate step.
# Section 7.19.3.


def bridge_mean(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """m_t = (1 - t) x_0 + t x_1, the mean of the bridge between the two ends."""
    return (1.0 - t) * x0 + t * x1


def bridge_std(t: torch.Tensor, epsilon: float = EPSILON) -> torch.Tensor:
    """sigma_t = sqrt(epsilon t (1 - t)), zero at both ends and largest at t = 1/2."""
    return torch.sqrt(epsilon * t * (1.0 - t))


def bridge_sample(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor, z: torch.Tensor,
                  epsilon: float = EPSILON) -> torch.Tensor:
    """x = m_t + sigma_t z, one draw from the bridge at time t."""
    return bridge_mean(x0, x1, t) + bridge_std(t, epsilon) * z


# %% 03. The two conditional quantities
#
# Both are conditional on the drawn pair, and neither depends on a quantity a network has to
# estimate first. Section 7.19.3.


def conditional_drift(x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """u_t(x | x_0, x_1) = (x_1 - x) / (1 - t), the remaining displacement over the remaining time."""
    return (x1 - x) / (1.0 - t)


def conditional_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
                      epsilon: float = EPSILON) -> torch.Tensor:
    """grad log p_t(x | x_0, x_1) = (m_t - x) / (epsilon t (1 - t)), displacement over variance."""
    return (bridge_mean(x0, x1, t) - x) / (epsilon * t * (1.0 - t))


# %% 04. From the bridge drift to a probability flow velocity
#
# We subtract half the noise variance times the score. That converts a diffusion drift into the
# velocity of the probability flow, and that flow has the same marginals with no noise at all.
# The second form is the straight-line velocity plus a coefficient times the displacement of the
# sample from the mean, with the dependence on epsilon cancelled. Section 7.19.3.


def flow_coefficient(t: torch.Tensor) -> torch.Tensor:
    """(1 - 2t) / (2 t (1 - t)), positive before t = 1/2 and negative after."""
    return (1.0 - 2.0 * t) / (2.0 * t * (1.0 - t))


def flow_velocity(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
                  epsilon: float = EPSILON) -> torch.Tensor:
    """u_t - (epsilon / 2) grad log p_t, the velocity the first network regresses onto."""
    return (conditional_drift(x1, x, t)
            - 0.5 * epsilon * conditional_score(x0, x1, x, t, epsilon))


def flow_velocity_collected(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor,
                            t: torch.Tensor) -> torch.Tensor:
    """(x_1 - x_0) + (1 - 2t) / (2 t (1 - t)) (x - m_t), the same velocity with delta collected."""
    return (x1 - x0) + flow_coefficient(t) * (x - bridge_mean(x0, x1, t))


def scaled_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
                 epsilon: float = EPSILON) -> torch.Tensor:
    """sigma_t grad log p_t = (m_t - x) / sigma_t = -z, with unit scale at every time."""
    return bridge_std(t, epsilon) * conditional_score(x0, x1, x, t, epsilon)


# %% 05. The loss and the sampler
#
# With lambda(t) = sigma_t the score target is -z, so no part of the schedule dominates the
# gradient. At sampling time we fix a noise level g(t) and integrate the two networks together,
# with g^2 = epsilon recovering the reference noise and g = 0 the probability flow.
# Section 7.19.3.


def sf2m_loss(velocity: torch.Tensor, score: torch.Tensor, velocity_target: torch.Tensor,
              score_target: torch.Tensor, t: torch.Tensor, epsilon: float = EPSILON
              ) -> torch.Tensor:
    """||v_theta - u^o_t||^2 + lambda(t)^2 ||s_theta - grad log p_t||^2, with lambda(t) = sigma_t."""
    weight = bridge_std(t, epsilon) ** 2  # lambda(t)^2, and it turns the score target into -z.
    return ((velocity - velocity_target) ** 2 + weight * (score - score_target) ** 2).mean()


def sampler_drift(velocity: torch.Tensor, score: torch.Tensor, noise_level: float
                  ) -> torch.Tensor:
    """v_theta + (1/2) g(t)^2 s_theta, the drift of the family of samplers between the two ends."""
    return velocity + 0.5 * noise_level ** 2 * score


# %% 06. The role of the coupling
#
# If we draw the pairs from the entropic plan at strength 2 epsilon, the learned equation at
# g^2 = epsilon is the Schrodinger bridge, and from any other coupling it is the Markov projection
# of that mixture of bridges. We compute the plan on a minibatch. Section 7.19.3.


def entropic_plan(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
                  epsilon: float = EPSILON, sweeps: int = 200) -> torch.Tensor:
    """The entropic plan for the squared cost at strength 2 epsilon, by Sinkhorn on exp(-C / 2 eps)."""
    C = (x[:, None] - y[None, :]) ** 2  # [n, m] the squared cost of the batch.
    K = torch.exp(-C / (2.0 * epsilon))  # The Gibbs kernel a Brownian reference produces.
    u, v = torch.ones_like(a), torch.ones_like(b)
    for _ in range(sweeps):
        u = a / (K @ v)
        v = b / (K.T @ u)
    return u[:, None] * K * v[None, :]  # [n, m].


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=101)  # --steps is the number of times on the grid.
    args = parser.parse_args()
    set_seed(args.seed)
    loud = not args.quiet
    x0 = torch.tensor(X0, dtype=DTYPE)
    x1 = torch.tensor(X1, dtype=DTYPE)
    z = torch.tensor(DRAW, dtype=DTYPE)

    worked = {}
    for value in (0.25, 0.75):
        t = torch.tensor(value, dtype=DTYPE)
        x = bridge_sample(x0, x1, t, z)
        worked[value] = {
            "mean": float(bridge_mean(x0, x1, t)),
            "std": float(bridge_std(t)),
            "sample": float(x),
            "drift": float(conditional_drift(x1, x, t)),
            "score": float(conditional_score(x0, x1, x, t)),
            "flow": float(flow_velocity(x0, x1, x, t)),
            "collected": float(flow_velocity_collected(x0, x1, x, t)),
            "scaled": float(scaled_score(x0, x1, x, t)),
            "coefficient": float(flow_coefficient(t)),
        }
        if loud:
            w = worked[value]
            print(f"t = {value}, z = {DRAW}")
            print(f"  m_t = {w['mean']:.4f}   sigma_t = {w['std']:.4f}   x = {w['sample']:.4f}")
            print(f"  u_t = {w['drift']:.4f}   grad log p_t = {w['score']:.4f}")
            print(f"  u^o_t = {w['flow']:.4f}   collected form {w['collected']:.4f}   "
                  f"coefficient {w['coefficient']:+.4f}")
            print(f"  sigma_t grad log p_t = {w['scaled']:.4f}, which is -z")

    grid = torch.linspace(0.0, 1.0, max(args.steps, 3), dtype=DTYPE)[1:-1]  # Interior times.
    samples = bridge_sample(x0, x1, grid, z)  # [levels] one bridge sample per time.
    scaled = scaled_score(x0, x1, samples, grid)
    scaled_gap = float((scaled + z).abs().max())  # The target is -z at every time.
    flow_targets = flow_velocity(x0, x1, samples, grid)
    score_targets = conditional_score(x0, x1, samples, grid)
    loss_exact = float(sf2m_loss(flow_targets, score_targets, flow_targets, score_targets, grid))
    loss_zero = float(sf2m_loss(torch.zeros_like(flow_targets), torch.zeros_like(score_targets),
                                flow_targets, score_targets, grid))  # An untrained pair of networks.
    stratified = torch.special.ndtri(
        (torch.arange(max(args.steps, 3), dtype=DTYPE) + 0.5) / max(args.steps, 3))
    half = torch.tensor(0.5, dtype=DTYPE)
    averaged = float(flow_velocity(x0, x1, bridge_sample(x0, x1, half, stratified), half).mean())
    symmetry = worked[0.25]["flow"] + worked[0.75]["flow"] - 2.0 * (X1 - X0)
    if loud:
        print(f"\nOn {len(grid)} interior times")
        print(f"  largest gap between sigma_t grad log p_t and -z: {scaled_gap:.2e}")
        print(f"  u^o averaged over the draw at t = 1/2: {averaged:.4f}, "
              f"against the straight-line velocity {X1 - X0:.0f}")
        print(f"  the two flow targets about the straight-line velocity: {symmetry:.2e}")
        print(f"  loss at the exact targets {loss_exact:.2e}, at a zero network {loss_zero:.4f}")

    x = torch.tensor([0.0, 2.0], dtype=DTYPE)  # The batch of Section 7.8.1.
    y = torch.tensor([1.0, 3.0], dtype=DTYPE)
    uniform = torch.full((2,), 0.5, dtype=DTYPE)
    C = (x[:, None] - y[None, :]) ** 2
    plan = entropic_plan(x, y, uniform, uniform)
    ordered_cost = float(C[0, 0] + C[1, 1])  # 1 + 1, the arrows of the exact pairing.
    crossed_cost = float(C[0, 1] + C[1, 0])  # 9 + 1, the arrows of the crossed pairing.
    if loud:
        print(f"\nThe two-point batch at epsilon = {EPSILON:g}")
        print(f"  exact transport pairs 0 with 1 and 2 with 3 at cost {ordered_cost:.0f}, "
              f"against {crossed_cost:.0f} for the crossed pairing")
        print(f"  entropic plan at strength {2 * EPSILON:g}: "
              f"{[[round(float(z_), 6) for z_ in row] for row in plan]}")
        print(f"  mass on the crossed pairing {float(plan[0, 1] + plan[1, 0]):.6f}")

    report = {
        "m_t at t = 0.25": worked[0.25]["mean"],
        "sigma_t at t = 0.25": worked[0.25]["std"],
        "bridge sample at t = 0.25": worked[0.25]["sample"],
        "u_t at t = 0.25": worked[0.25]["drift"],
        "grad log p_t at t = 0.25": worked[0.25]["score"],
        "u^o_t at t = 0.25": worked[0.25]["flow"],
        "coefficient at t = 0.25": worked[0.25]["coefficient"],
        "sigma_t grad log p_t at t = 0.25": worked[0.25]["scaled"],
        "bridge sample at t = 0.75": worked[0.75]["sample"],
        "u_t at t = 0.75": worked[0.75]["drift"],
        "u^o_t at t = 0.75": worked[0.75]["flow"],
        "coefficient at t = 0.75": worked[0.75]["coefficient"],
        "the two targets, sum minus 2 u": symmetry,
        "largest gap of the scaled score": scaled_gap,
        "u^o averaged over the draw": averaged,
        "loss at the exact targets": loss_exact,
        "loss at a zero network": loss_zero,
        "sampler drift at g^2 = epsilon": float(sampler_drift(
            torch.tensor(worked[0.25]["flow"], dtype=DTYPE),
            torch.tensor(worked[0.25]["score"], dtype=DTYPE), EPSILON ** 0.5)),
        "sampler drift at g = 0": worked[0.25]["flow"],
        "ordered pairing of the batch": ordered_cost,
        "crossed pairing of the batch": crossed_cost,
        "entropic plan, ordered mass": float(plan[0, 0] + plan[1, 1]),
        "entropic plan, crossed mass": float(plan[0, 1] + plan[1, 0]),
    }
    print_report("Simulation-Free Score and Flow Matching, Chapter 7 Section 7.19.3", report)
    save_run(args.out, config=vars(args), report=report)  # Written only with --out.


if __name__ == "__main__":
    main()

	"""Simulation-Free Score and Flow Matching: two closed-form targets at one bridge sample.

	With the entropic plan computed on a minibatch, the bridge is a mixture of Brownian bridges between
	pairs drawn from that plan, and a Brownian bridge at a time t is Gaussian. One normal draw gives a
	sample, and at that sample both the bridge drift and the score of the bridge marginal are available
	without simulating a path. We subtract half the noise variance times the score. That turns the
	drift into the velocity of the probability flow, and its dependence on epsilon cancels.

	Chapter 7, Section 7.19.3 of the notes. Runs on the endpoint pair x_0 = 0, x_1 = 2 at epsilon = 1,
	and on the two-point batch of Section 7.8.1. Prints the two regression targets at t = 0.25 and at
	t = 0.75, the scaled score at every time on the grid, and the entropic plan of the batch.

	python lecture_7/sf2m.py # 101 times on the grid
	python lecture_7/sf2m.py --steps 20 # a coarse time grid, for a first look
	"""

	# %% 01. Imports and configuration
	from __future__ import annotations

	import torch

	from cis6270.runner import common_parser, print_report, save_run, set_seed

	DTYPE = torch.float64 # The chapter quotes four figures of both regression targets.

	X0, X1 = 0.0, 2.0 # The endpoint pair the chapter works through.
	EPSILON = 1.0 # The reference noise variance per unit time.
	DRAW = 0.4 # The normal draw z the chapter uses at both times.


	# %% 02. One sample from the bridge
	#
	# The Brownian bridge at time t is Gaussian with mean the straight interpolation and variance
	# epsilon t (1 - t), so a sample costs one normal draw and no intermediate step.
	# Section 7.19.3.


	def bridge_mean(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
	"""m_t = (1 - t) x_0 + t x_1, the mean of the bridge between the two ends."""
	return (1.0 - t) * x0 + t * x1


	def bridge_std(t: torch.Tensor, epsilon: float = EPSILON) -> torch.Tensor:
	"""sigma_t = sqrt(epsilon t (1 - t)), zero at both ends and largest at t = 1/2."""
	return torch.sqrt(epsilon * t * (1.0 - t))


	def bridge_sample(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor, z: torch.Tensor,
	epsilon: float = EPSILON) -> torch.Tensor:
	"""x = m_t + sigma_t z, one draw from the bridge at time t."""
	return bridge_mean(x0, x1, t) + bridge_std(t, epsilon) * z


	# %% 03. The two conditional quantities
	#
	# Both are conditional on the drawn pair, and neither depends on a quantity a network has to
	# estimate first. Section 7.19.3.


	def conditional_drift(x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
	"""u_t(x \| x_0, x_1) = (x_1 - x) / (1 - t), the remaining displacement over the remaining time."""
	return (x1 - x) / (1.0 - t)


	def conditional_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
	epsilon: float = EPSILON) -> torch.Tensor:
	"""grad log p_t(x \| x_0, x_1) = (m_t - x) / (epsilon t (1 - t)), displacement over variance."""
	return (bridge_mean(x0, x1, t) - x) / (epsilon * t * (1.0 - t))


	# %% 04. From the bridge drift to a probability flow velocity
	#
	# We subtract half the noise variance times the score. That converts a diffusion drift into the
	# velocity of the probability flow, and that flow has the same marginals with no noise at all.
	# The second form is the straight-line velocity plus a coefficient times the displacement of the
	# sample from the mean, with the dependence on epsilon cancelled. Section 7.19.3.


	def flow_coefficient(t: torch.Tensor) -> torch.Tensor:
	"""(1 - 2t) / (2 t (1 - t)), positive before t = 1/2 and negative after."""
	return (1.0 - 2.0 * t) / (2.0 * t * (1.0 - t))


	def flow_velocity(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
	epsilon: float = EPSILON) -> torch.Tensor:
	"""u_t - (epsilon / 2) grad log p_t, the velocity the first network regresses onto."""
	return (conditional_drift(x1, x, t)
	- 0.5 * epsilon * conditional_score(x0, x1, x, t, epsilon))


	def flow_velocity_collected(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor,
	t: torch.Tensor) -> torch.Tensor:
	"""(x_1 - x_0) + (1 - 2t) / (2 t (1 - t)) (x - m_t), the same velocity with delta collected."""
	return (x1 - x0) + flow_coefficient(t) * (x - bridge_mean(x0, x1, t))


	def scaled_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor,
	epsilon: float = EPSILON) -> torch.Tensor:
	"""sigma_t grad log p_t = (m_t - x) / sigma_t = -z, with unit scale at every time."""
	return bridge_std(t, epsilon) * conditional_score(x0, x1, x, t, epsilon)


	# %% 05. The loss and the sampler
	#
	# With lambda(t) = sigma_t the score target is -z, so no part of the schedule dominates the
	# gradient. At sampling time we fix a noise level g(t) and integrate the two networks together,
	# with g^2 = epsilon recovering the reference noise and g = 0 the probability flow.
	# Section 7.19.3.


	def sf2m_loss(velocity: torch.Tensor, score: torch.Tensor, velocity_target: torch.Tensor,
	score_target: torch.Tensor, t: torch.Tensor, epsilon: float = EPSILON
	) -> torch.Tensor:
	"""\|\|v_theta - u^o_t\|\|^2 + lambda(t)^2 \|\|s_theta - grad log p_t\|\|^2, with lambda(t) = sigma_t."""
	weight = bridge_std(t, epsilon) ** 2 # lambda(t)^2, and it turns the score target into -z.
	return ((velocity - velocity_target) ** 2 + weight * (score - score_target) ** 2).mean()


	def sampler_drift(velocity: torch.Tensor, score: torch.Tensor, noise_level: float
	) -> torch.Tensor:
	"""v_theta + (1/2) g(t)^2 s_theta, the drift of the family of samplers between the two ends."""
	return velocity + 0.5 * noise_level ** 2 * score


	# %% 06. The role of the coupling
	#
	# If we draw the pairs from the entropic plan at strength 2 epsilon, the learned equation at
	# g^2 = epsilon is the Schrodinger bridge, and from any other coupling it is the Markov projection
	# of that mixture of bridges. We compute the plan on a minibatch. Section 7.19.3.


	def entropic_plan(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor,
	epsilon: float = EPSILON, sweeps: int = 200) -> torch.Tensor:
	"""The entropic plan for the squared cost at strength 2 epsilon, by Sinkhorn on exp(-C / 2 eps)."""
	C = (x[:, None] - y[None, :]) ** 2 # [n, m] the squared cost of the batch.
	K = torch.exp(-C / (2.0 * epsilon)) # The Gibbs kernel a Brownian reference produces.
	u, v = torch.ones_like(a), torch.ones_like(b)
	for _ in range(sweeps):
	u = a / (K @ v)
	v = b / (K.T @ u)
	return u[:, None] * K * v[None, :] # [n, m].


	# %% 07. The report
	def main() -> None:
	parser = common_parser(__doc__.split("\n")[0])
	parser.set_defaults(steps=101) # --steps is the number of times on the grid.
	args = parser.parse_args()
	set_seed(args.seed)
	loud = not args.quiet
	x0 = torch.tensor(X0, dtype=DTYPE)
	x1 = torch.tensor(X1, dtype=DTYPE)
	z = torch.tensor(DRAW, dtype=DTYPE)

	worked = {}
	for value in (0.25, 0.75):
	t = torch.tensor(value, dtype=DTYPE)
	x = bridge_sample(x0, x1, t, z)
	worked[value] = {
	"mean": float(bridge_mean(x0, x1, t)),
	"std": float(bridge_std(t)),
	"sample": float(x),
	"drift": float(conditional_drift(x1, x, t)),
	"score": float(conditional_score(x0, x1, x, t)),
	"flow": float(flow_velocity(x0, x1, x, t)),
	"collected": float(flow_velocity_collected(x0, x1, x, t)),
	"scaled": float(scaled_score(x0, x1, x, t)),
	"coefficient": float(flow_coefficient(t)),
	}
	if loud:
	w = worked[value]
	print(f"t = {value}, z = {DRAW}")
	print(f" m_t = {w['mean']:.4f} sigma_t = {w['std']:.4f} x = {w['sample']:.4f}")
	print(f" u_t = {w['drift']:.4f} grad log p_t = {w['score']:.4f}")
	print(f" u^o_t = {w['flow']:.4f} collected form {w['collected']:.4f} "
	f"coefficient {w['coefficient']:+.4f}")
	print(f" sigma_t grad log p_t = {w['scaled']:.4f}, which is -z")

	grid = torch.linspace(0.0, 1.0, max(args.steps, 3), dtype=DTYPE)[1:-1] # Interior times.
	samples = bridge_sample(x0, x1, grid, z) # [levels] one bridge sample per time.
	scaled = scaled_score(x0, x1, samples, grid)
	scaled_gap = float((scaled + z).abs().max()) # The target is -z at every time.
	flow_targets = flow_velocity(x0, x1, samples, grid)
	score_targets = conditional_score(x0, x1, samples, grid)
	loss_exact = float(sf2m_loss(flow_targets, score_targets, flow_targets, score_targets, grid))
	loss_zero = float(sf2m_loss(torch.zeros_like(flow_targets), torch.zeros_like(score_targets),
	flow_targets, score_targets, grid)) # An untrained pair of networks.
	stratified = torch.special.ndtri(
	(torch.arange(max(args.steps, 3), dtype=DTYPE) + 0.5) / max(args.steps, 3))
	half = torch.tensor(0.5, dtype=DTYPE)
	averaged = float(flow_velocity(x0, x1, bridge_sample(x0, x1, half, stratified), half).mean())
	symmetry = worked[0.25]["flow"] + worked[0.75]["flow"] - 2.0 * (X1 - X0)
	if loud:
	print(f"\nOn {len(grid)} interior times")
	print(f" largest gap between sigma_t grad log p_t and -z: {scaled_gap:.2e}")
	print(f" u^o averaged over the draw at t = 1/2: {averaged:.4f}, "
	f"against the straight-line velocity {X1 - X0:.0f}")
	print(f" the two flow targets about the straight-line velocity: {symmetry:.2e}")
	print(f" loss at the exact targets {loss_exact:.2e}, at a zero network {loss_zero:.4f}")

	x = torch.tensor([0.0, 2.0], dtype=DTYPE) # The batch of Section 7.8.1.
	y = torch.tensor([1.0, 3.0], dtype=DTYPE)
	uniform = torch.full((2,), 0.5, dtype=DTYPE)
	C = (x[:, None] - y[None, :]) ** 2
	plan = entropic_plan(x, y, uniform, uniform)
	ordered_cost = float(C[0, 0] + C[1, 1]) # 1 + 1, the arrows of the exact pairing.
	crossed_cost = float(C[0, 1] + C[1, 0]) # 9 + 1, the arrows of the crossed pairing.
	if loud:
	print(f"\nThe two-point batch at epsilon = {EPSILON:g}")
	print(f" exact transport pairs 0 with 1 and 2 with 3 at cost {ordered_cost:.0f}, "
	f"against {crossed_cost:.0f} for the crossed pairing")
	print(f" entropic plan at strength {2 * EPSILON:g}: "
	f"{[[round(float(z_), 6) for z_ in row] for row in plan]}")
	print(f" mass on the crossed pairing {float(plan[0, 1] + plan[1, 0]):.6f}")

	report = {
	"m_t at t = 0.25": worked[0.25]["mean"],
	"sigma_t at t = 0.25": worked[0.25]["std"],
	"bridge sample at t = 0.25": worked[0.25]["sample"],
	"u_t at t = 0.25": worked[0.25]["drift"],
	"grad log p_t at t = 0.25": worked[0.25]["score"],
	"u^o_t at t = 0.25": worked[0.25]["flow"],
	"coefficient at t = 0.25": worked[0.25]["coefficient"],
	"sigma_t grad log p_t at t = 0.25": worked[0.25]["scaled"],
	"bridge sample at t = 0.75": worked[0.75]["sample"],
	"u_t at t = 0.75": worked[0.75]["drift"],
	"u^o_t at t = 0.75": worked[0.75]["flow"],
	"coefficient at t = 0.75": worked[0.75]["coefficient"],
	"the two targets, sum minus 2 u": symmetry,
	"largest gap of the scaled score": scaled_gap,
	"u^o averaged over the draw": averaged,
	"loss at the exact targets": loss_exact,
	"loss at a zero network": loss_zero,
	"sampler drift at g^2 = epsilon": float(sampler_drift(
	torch.tensor(worked[0.25]["flow"], dtype=DTYPE),
	torch.tensor(worked[0.25]["score"], dtype=DTYPE), EPSILON ** 0.5)),
	"sampler drift at g = 0": worked[0.25]["flow"],
	"ordered pairing of the batch": ordered_cost,
	"crossed pairing of the batch": crossed_cost,
	"entropic plan, ordered mass": float(plan[0, 0] + plan[1, 1]),
	"entropic plan, crossed mass": float(plan[0, 1] + plan[1, 0]),
	}
	print_report("Simulation-Free Score and Flow Matching, Chapter 7 Section 7.19.3", report)
	save_run(args.out, config=vars(args), report=report) # Written only with --out.


	if __name__ == "__main__":
	main()