flow.unrolled_populations

flow.unrolled_populations(step_fn, theta, f0, steps)

The same march with the tape left on, for cross-checking the adjoint.

Reverse mode through this stores all steps intermediate states. It exists to prove :func:steady_populations right at a size where the tape still fits, and to measure what it costs – not to be used.

Parameters

Name	Type	Description	Default
step_fn	Callable	(f, theta) -> f.	required
theta	Any	The differentiable parameters.	required
f0	jax.Array	Initial populations.	required
steps	int	Exactly how many steps to run.	required

Returns

Name	Type	Description
	jax.Array	The populations after steps steps.