Rename backend.control_flow to backend.flow to shorten

Author: pnkraemer

probdiffeq/backend/control_flow.py probdiffeq/backend/flow.pyprobdiffeq/backend/control_flow.py renamed to probdiffeq/backend/flow.py

@@ -7,15 +7,7 @@
 See the tutorials for example use cases.
 """
 
-from probdiffeq.backend import (
-    containers,
-    control_flow,
-    func,
-    linalg,
-    np,
-    tree,
-    warnings,
-)
+from probdiffeq.backend import containers, flow, func, linalg, np, tree, warnings
 from probdiffeq.backend.typing import Any, Array, Callable, Generic, Protocol, TypeVar
 
 T_contra = TypeVar("T_contra", contravariant=True)
@@ -147,7 +139,7 @@ def solve_adaptive_terminal_values(
     errorest,
     control: Control | None = None,
     clip_dt: bool = False,
-    while_loop: Callable = control_flow.while_loop,
+    while_loop: Callable = flow.while_loop,
 ) -> Callable[..., Solution]:
     """Simulate the terminal values of an initial value problem."""
     # Turn off warnings because any solver goes for terminal values
@@ -178,7 +170,7 @@ def solve_adaptive_save_at(
     errorest,
     control: Control | None = None,
     clip_dt: bool = False,
-    while_loop: Callable = control_flow.while_loop,
+    while_loop: Callable = flow.while_loop,
     warn=True,
 ) -> Callable[..., Solution]:
     r"""Solve an initial value problem and return the solution at a pre-determined grid.
@@ -261,7 +253,7 @@ def body_fun(state: AdvanceState) -> AdvanceState:
         # Advance to one checkpoint after the other
         init = (solution0, state)
         xs = save_at[1:]
-        (_solution, _state), solution = control_flow.scan(
+        (_solution, _state), solution = flow.scan(
             advance, init=init, xs=xs, reverse=False
         )
 
@@ -281,7 +273,7 @@ def body_fn(s, dt):
 
         t0 = grid[0]
         state0 = solver.init(t=t0, u=u)
-        _, result = control_flow.scan(body_fn, init=state0, xs=np.diff(grid))
+        _, result = flow.scan(body_fn, init=state0, xs=np.diff(grid))
 
         return solver.userfriendly_output(solution0=state0, solution=result)
 
@@ -426,14 +418,14 @@ def loop(
         # If t1 is in the future, enter the rejection loop (otherwise do nothing)
         is_before_t1 = state0.step_from.t + eps < t1
         args = (state0, t1, atol, rtol, damp)
-        state = control_flow.cond(is_before_t1, self.step, lambda s: s[0], args)
+        state = flow.cond(is_before_t1, self.step, lambda s: s[0], args)
 
         # Interpolate
         is_before_t1 = state.step_from.t + eps < t1
         is_after_t1 = state.step_from.t > t1 + eps
         branch_idx = np.where(is_before_t1, 0, np.where(is_after_t1, 1, 2))
         options = (self.interp_skip, self.interp_beyond_t1, self.interp_at_t1)
-        return control_flow.switch(branch_idx, options, (state, t1))
+        return flow.switch(branch_idx, options, (state, t1))
 
     def step(self, s_and_t1_and_tols_and_damp):
         """Do a rejection-loop step.

probdiffeq/ivpsolve.py

@@ -3,16 +3,7 @@
 See the tutorials for example use cases.
 """
 
-from probdiffeq.backend import (
-    containers,
-    control_flow,
-    func,
-    linalg,
-    np,
-    random,
-    special,
-    tree,
-)
+from probdiffeq.backend import containers, flow, func, linalg, np, random, special, tree
 from probdiffeq.backend.typing import (
     Any,
     Array,
@@ -324,7 +315,7 @@ def step(x, cond):
             return extrapolated, extrapolated
 
         init, xs = markov_seq.marginal, markov_seq.conditional
-        _, marginals = control_flow.scan(step, init=init, xs=xs, reverse=reverse)
+        _, marginals = flow.scan(step, init=init, xs=xs, reverse=reverse)
 
         if reverse:
             # Append the terminal marginal to the computed ones
@@ -386,9 +377,7 @@ def body_fun(samp_prev, conditionals_and_base_samples):
 
         # Loop over backward models and the remaining base samples
         xs = (markov_seq.conditional, base_sample_body)
-        _, samples = control_flow.scan(
-            body_fun, init=init_sample, xs=xs, reverse=reverse
-        )
+        _, samples = flow.scan(body_fun, init=init_sample, xs=xs, reverse=reverse)
 
         if reverse:
             samples = np.concatenate([samples, init_sample[None, ...]])

probdiffeq/probdiffeq.py

@@ -7,7 +7,7 @@
 See the tutorials for example use cases.
 """
 
-from probdiffeq.backend import control_flow, func, np, ode, tree
+from probdiffeq.backend import flow, func, np, ode, tree
 from probdiffeq.backend.typing import Array, ArrayLike, Callable, Sequence
 from probdiffeq.util import filter_util
 
@@ -126,9 +126,7 @@ def body(tcoeffs, _):
         return taylor_coeffs
 
     # Compute all coefficients with scan().
-    taylor_coeffs, _ = control_flow.scan(
-        body, init=taylor_coeffs, xs=None, length=num - 1
-    )
+    taylor_coeffs, _ = flow.scan(body, init=taylor_coeffs, xs=None, length=num - 1)
     return taylor_coeffs
 
 
@@ -329,7 +327,7 @@ def body_fun(cs_padded, i_and_fx_i):
         cs_padded = np.stack(cs_padded)
 
         xs = [np.arange(0, len(fx[deg : 2 * deg])), fx[deg : 2 * deg]]
-        cs_padded, _ = control_flow.scan(body_fun, xs=xs, init=cs_padded)
+        cs_padded, _ = flow.scan(body_fun, xs=xs, init=cs_padded)
 
         taylor_coefficients.extend(cs_padded)
         return taylor_coefficients

probdiffeq/taylor.py

@@ -20,7 +20,7 @@
 
 """
 
-from probdiffeq.backend import control_flow, linalg, np, tree
+from probdiffeq.backend import flow, linalg, np, tree
 
 
 def revert_conditional_noisefree(R_X_F, R_X):
@@ -187,7 +187,7 @@ def f_body(idx, f):
         )
         return f.at[idx].set(val)
 
-    f = control_flow.fori_loop(1, n, f_body, f)
+    f = flow.fori_loop(1, n, f_body, f)
     f = 1.0 / f
 
     U = np.eye(n)
@@ -204,10 +204,10 @@ def inner_body(k, g):
             newval = (g[i + 1] / denom) * factor
             return g.at[i].set(newval)
 
-        g = control_flow.fori_loop(0, j_idx, inner_body, g)
+        g = flow.fori_loop(0, j_idx, inner_body, g)
         return U.at[:, j_idx].set(g)
 
-    U = control_flow.fori_loop(1, n, body_j, U)
+    U = flow.fori_loop(1, n, body_j, U)
 
     # scale columns: U = U .* (dr * f_row)
     U = U * (dr[:, None] * f[None, :])

probdiffeq/util/cholesky_util.py

@@ -3,7 +3,7 @@
 Mostly **discrete** filtering and smoothing.
 """
 
-from probdiffeq.backend import containers, control_flow, tree
+from probdiffeq.backend import containers, flow, tree
 from probdiffeq.backend.typing import Any
 
 
@@ -19,7 +19,7 @@ def estimate_fwd(data, /, init, prior_transitions, observation_model, *, estimat
     idx_or_slice = slice(1, len(data), 1)
     information = _select((data, observation_model), idx_or_slice=idx_or_slice)
     xs = (prior_transitions, *information)
-    return control_flow.scan(step, init=init, xs=xs, reverse=False)
+    return flow.scan(step, init=init, xs=xs, reverse=False)
 
 
 def estimate_rev(data, /, init, prior_transitions, observation_model, *, estimator):
@@ -33,7 +33,7 @@ def estimate_rev(data, /, init, prior_transitions, observation_model, *, estimat
     # Scan over the remaining data points
     information = _select((data, observation_model), idx_or_slice=slice(0, -1, 1))
     xs = (prior_transitions, *information)
-    return control_flow.scan(step, init=init, xs=xs, reverse=True)
+    return flow.scan(step, init=init, xs=xs, reverse=True)
 
 
 def fixedpointsmoother_precon(*, ssm):

probdiffeq/util/filter_util.py

@@ -1,7 +1,7 @@
 """Test utilities."""
 
 import probdiffeq.ivpsolve
-from probdiffeq.backend import control_flow, func, tree, warnings
+from probdiffeq.backend import flow, func, tree, warnings
 from probdiffeq.backend.typing import TypeVar
 
 T = TypeVar("T")
@@ -35,7 +35,7 @@ def solve_adaptive_save_every_step(solver, errorest, control=None, clip_dt=False
         # We do not expose this option to the user
         # because we do not want to suggest that this function
         # uses meaningful looping to begin with.
-        while_loop=control_flow.while_loop,
+        while_loop=flow.while_loop,
     )
 
     def solve(