flow.steady.steady_populations
flow.steady.steady_populations(step_fn, options, theta, f0)The converged populations, differentiated through the fixed point.
Forward, this is :func:iterate_to_steady. Backward, it is the adjoint system of the converged equation – the initial guess f0 therefore receives a zero cotangent, which is the honest answer: move the guess and the converged state does not move.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| step_fn | Callable | (f, theta) -> f, one lattice time. Static; passed as a nondiff_argnum, so a fresh closure per call defeats caching. |
required |
| options | SteadyOptions | Convergence settings, static and hashable. | required |
| theta | Any | The differentiable parameters (any pytree step_fn accepts). |
required |
| f0 | jax.Array | Initial populations. | required |
Returns
| Name | Type | Description |
|---|---|---|
| jax.Array | The converged populations, (19, NX, NY, NZ). |