flow.step_for
flow.step_for(config)The memoised one-step closure for a configuration.
Cached on config so that repeated solves reuse one compiled scan; see this module’s docstring for why identity matters here.
Parameters
| Name | Type | Description | Default |
|---|---|---|---|
| config | FlowConfig | The flow configuration. | required |
Returns
| Name | Type | Description |
|---|---|---|
| Callable[[jax.Array, dict], jax.Array] | step_fn(f, theta) -> f where theta has keys "chi" and |
|
| Callable[[jax.Array, dict], jax.Array] | "inlet_velocity". |