diff --git a/cuthbert/README.md b/cuthbert/README.md
index 418139f7..d9b8741c 100644
--- a/cuthbert/README.md
+++ b/cuthbert/README.md
@@ -10,18 +10,18 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-init_model_inputs = ...
 filter_model_inputs = ...
 
 # Load inference method
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
-    get_init_params=get_init_params,
+    m0=m0,
+    chol_P0=chol_P0,
     get_dynamics_params=get_dynamics_params,
     get_observation_params=get_observation_params,
 )   # build_filter function takes all inference-specific arguments, swap this out for different inference methods.
 
 # Online inference
-state = kalman_filter.init_prepare(init_model_inputs)
+state = kalman_filter.init_prepare()
 
 for t in range(T):
     model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
@@ -34,7 +34,7 @@ Or for offline inference:
 ```python
 kalman_smoother = cuthbert.gaussian.kalman.build_smoother(get_dynamics_params)
 
-init_state = kalman_filter.init_prepare(init_model_inputs)
+init_state = kalman_filter.init_prepare()
 filter_states = cuthbert.filter(kalman_filter, filter_model_inputs, init_state)
 smoother_states = cuthbert.smoother(kalman_smoother, filter_states, filter_model_inputs)
 ```
diff --git a/cuthbert/discrete/filter.py b/cuthbert/discrete/filter.py
index 0548f28d..37d7114b 100644
--- a/cuthbert/discrete/filter.py
+++ b/cuthbert/discrete/filter.py
@@ -13,7 +13,6 @@
 from jax import tree
 
 from cuthbert.discrete.types import (
-    GetInitDist,
     GetObsLogLikelihoods,
     GetTransitionMatrix,
 )
@@ -43,14 +42,14 @@ def log_normalizing_constant(self) -> Array:
 
 
 def build_filter(
-    get_init_dist: GetInitDist,
+    init_dist: Array,
     get_trans_matrix: GetTransitionMatrix,
     get_obs_lls: GetObsLogLikelihoods,
 ) -> Filter:
     r"""Builds a filter object for discrete hidden Markov models.
 
     Args:
-        get_init_dist: Function to get initial state probabilities $m_i = p(x_0 = i)$.
+        init_dist: Array of initial state probabilities $m_i = p(x_0 = i)$.
         get_trans_matrix: Function to get the transition matrix $A_{ij} = p(x_t = j \mid x_{t-1} = i)$.
         get_obs_lls: Function to get observation log likelihoods $b_i = \log p(y_t | x_t = i)$.
 
@@ -58,7 +57,7 @@ def build_filter(
         Filter object. Suitable for associative scan.
     """
     return Filter(
-        init_prepare=partial(init_prepare, get_init_dist=get_init_dist),
+        init_prepare=partial(init_prepare, init_dist=init_dist),
         filter_prepare=partial(
             filter_prepare, get_trans_matrix=get_trans_matrix, get_obs_lls=get_obs_lls
         ),
@@ -67,27 +66,23 @@ def build_filter(
     )
 
 
-def init_prepare(
-    model_inputs: ArrayTreeLike, get_init_dist: GetInitDist, key: KeyArray | None = None
-) -> DiscreteFilterState:
+def init_prepare(init_dist: Array, key: KeyArray | None = None) -> DiscreteFilterState:
     """Prepare the initial state for the filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_dist: Function to get initial state probabilities m_i = p(x_0 = i).
+        init_dist: Array of initial state probabilities m_i = p(x_0 = i).
         key: JAX random key - not used.
 
     Returns:
         Prepared state for the filter.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    init_dist = get_init_dist(model_inputs)
+    init_dist = jnp.asarray(init_dist)
     N = init_dist.shape[-1]
     f = init_dist[..., None, :] * jnp.ones((N, N))
     # repeat init_dist N times onto newly added penultimate axis
     log_g = jnp.zeros_like(init_dist)
     return DiscreteFilterState(
-        elem=filtering.FilterScanElement(f, log_g), model_inputs=model_inputs
+        elem=filtering.FilterScanElement(f, log_g), model_inputs=None
     )
 
 
diff --git a/cuthbert/discrete/types.py b/cuthbert/discrete/types.py
index 7f1b64b4..694c0c47 100644
--- a/cuthbert/discrete/types.py
+++ b/cuthbert/discrete/types.py
@@ -5,22 +5,6 @@
 from cuthbertlib.types import Array, ArrayTreeLike
 
 
-class GetInitDist(Protocol):
-    """Protocol for specifying the initial distribution."""
-
-    def __call__(self, model_inputs: ArrayTreeLike) -> Array:
-        """Get the initial distribution.
-
-        Args:
-            model_inputs: Model inputs.
-
-        Returns:
-            An array $m$ of shape (N,) where N is the number of states,
-                with $m_i = p(x_0 = i)$.
-        """
-        ...
-
-
 class GetTransitionMatrix(Protocol):
     """Protocol for specifying the transition matrix."""
 
diff --git a/cuthbert/ensemble_kalman/ensemble_kalman_filter.py b/cuthbert/ensemble_kalman/ensemble_kalman_filter.py
index 0c67a883..f89787d3 100644
--- a/cuthbert/ensemble_kalman/ensemble_kalman_filter.py
+++ b/cuthbert/ensemble_kalman/ensemble_kalman_filter.py
@@ -81,7 +81,8 @@ def build_filter(
     """Builds an Ensemble Kalman Filter object.
 
     Args:
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, generates a single sample from
+            the initial distribution.
         get_dynamics: Function to get dynamics function (x_t, key) -> x_{t+1} ~ p(x_{t+1} | x_t) from model inputs.
         get_observations: Function to get observation function, chol_R, and y from model inputs.
         n_particles: Number of particles.
@@ -134,7 +135,6 @@ def build_filter(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
     init_sample: InitSample,
     n_particles: int,
     store_predicted_ensemble: bool = False,
@@ -143,8 +143,8 @@ def init_prepare(
     """Prepare the initial state for the EnKF.
 
     Args:
-        model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, generates a single sample from
+            the initial distribution.
         n_particles: Number of particles.
         store_predicted_ensemble: Whether to store incoming forecast ensembles.
         key: JAX random key.
@@ -155,19 +155,18 @@ def init_prepare(
     Raises:
         ValueError: If key is None.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Sample ensemble from initial distribution
     keys = random.split(key, n_particles)
-    ensemble = jax.vmap(init_sample, (0, None))(keys, model_inputs)
+    ensemble = jax.vmap(init_sample)(keys)
     predicted_ensemble = dummy_tree_like(ensemble) if store_predicted_ensemble else None
 
     return EnKFState(
         key=key,
         ensemble=ensemble,
-        model_inputs=model_inputs,
+        model_inputs=None,
         log_normalizing_constant=jnp.array(0.0),
         predicted_ensemble=predicted_ensemble,
     )
@@ -184,7 +183,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution from key and model inputs.
+        init_sample: Function of a JAX random key only, sampling from the initial
+            distribution. Bind initial parameters when building the filter.
         n_particles: Number of particles.
         store_predicted_ensemble: Whether to store incoming forecast ensembles.
         key: JAX random key.
@@ -200,10 +200,11 @@ def filter_prepare(
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Infer state shape from init_sample
-    dummy_particle = jax.eval_shape(init_sample, key, model_inputs)
+    dummy_particle = jax.eval_shape(init_sample, key)
     x_dim = dummy_particle.shape[0]
-    ensemble = jnp.empty((n_particles, x_dim))
-    ensemble = dummy_tree_like(ensemble)
+    ensemble = dummy_tree_like(
+        jax.ShapeDtypeStruct((n_particles, x_dim), dummy_particle.dtype)
+    )
     predicted_ensemble = ensemble if store_predicted_ensemble else None
 
     return EnKFState(
diff --git a/cuthbert/ensemble_kalman/types.py b/cuthbert/ensemble_kalman/types.py
index 51848fd2..fd4ee70b 100644
--- a/cuthbert/ensemble_kalman/types.py
+++ b/cuthbert/ensemble_kalman/types.py
@@ -12,12 +12,11 @@
 class InitSample(Protocol):
     """Protocol for sampling from the initial distribution."""
 
-    def __call__(self, key: KeyArray, model_inputs: ArrayTreeLike) -> Array:
+    def __call__(self, key: KeyArray) -> Array:
         """Sample from the initial distribution.
 
         Args:
             key: JAX PRNG key.
-            model_inputs: Model inputs.
 
         Returns:
             Sample from the initial distribution, shape (x_dim,).
diff --git a/cuthbert/factorial/README.md b/cuthbert/factorial/README.md
index 898df3b4..7b3a2aa1 100644
--- a/cuthbert/factorial/README.md
+++ b/cuthbert/factorial/README.md
@@ -47,7 +47,6 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-init_model_inputs = ...
 filter_model_inputs = ...
 
 # Define function to extract the factorial indices from model inputs
@@ -59,14 +58,15 @@ factorializer = cuthbert.factorial.gaussian.build_factorializer(get_factorial_in
 
 # Load inference method, with parameter extraction functions defined for factorial inference
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
-    get_init_params=get_init_params,  # Init specified to generate factorial state
+    m0=m0,   # Shape (F, d) i.e. includes factorial dimension
+    chol_P0=chol_P0,  # Shape (F, d, d)
     get_dynamics_params=get_dynamics_params,  # Dynamics specified to act on joint local state
     get_observation_params=get_observation_params,  # Observation specified to act on joint local state
 )
 
 # Online inference
-factorial_state = kalman_filter.init_prepare(init_model_inputs)
-factorial_state = factorializer.factorialize_init_state(factorial_state, init_model_inputs)
+factorial_state = kalman_filter.init_prepare()
+factorial_state = factorializer.factorialize_init_state(factorial_state)
 
 for t in range(T):
     model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
@@ -82,8 +82,8 @@ You can also use `cuthbert.factorial.filter` for convenient offline filtering.
 Note that associative/parallel filtering is not supported for factorial filtering.
 
 ```python
-init_factorial_state = kalman_filter.init_prepare(init_model_inputs)
-init_factorial_state = factorializer.factorialize_init_state(init_factorial_state, init_model_inputs)
+init_factorial_state = kalman_filter.init_prepare()
+init_factorial_state = factorializer.factorialize_init_state(init_factorial_state)
 
 local_filter_states, final_factorial_state = (
     cuthbert.factorial.filter(
diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py
index 4ab62f81..63c32753 100644
--- a/cuthbert/factorial/filtering.py
+++ b/cuthbert/factorial/filtering.py
@@ -6,6 +6,7 @@
 
 from cuthbert.factorial.types import Factorializer
 from cuthbert.inference import Filter
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -61,6 +62,7 @@ def filter(
     """
     T = tree.leaves(model_inputs)[0].shape[0]
     prepare_keys = _prepare_keys(key, T)
+    init_state = init_state._replace(model_inputs=dummy_leading_element(model_inputs))
 
     def body_local(prev_factorial_state, prep_inp_and_k):
         prep_inp, k = prep_inp_and_k
diff --git a/cuthbert/factorial/smc.py b/cuthbert/factorial/smc.py
index ee2d37df..cc73697a 100644
--- a/cuthbert/factorial/smc.py
+++ b/cuthbert/factorial/smc.py
@@ -14,7 +14,7 @@
 from cuthbert.smc.marginal_particle_filter import MarginalParticleFilterState
 from cuthbert.smc.particle_filter import ParticleFilterState
 from cuthbertlib.resampling import Resampling, ess_decorator
-from cuthbertlib.types import Array, ArrayLike, ArrayTreeLike
+from cuthbertlib.types import Array, ArrayLike
 
 GeneralParticleFilterState = TypeVar(
     "GeneralParticleFilterState", ParticleFilterState, MarginalParticleFilterState
@@ -58,7 +58,7 @@ def build_factorializer(
 
 
 def factorialize_init_state(
-    init_state: GeneralParticleFilterState, model_inputs: ArrayTreeLike
+    init_state: GeneralParticleFilterState,
 ) -> GeneralParticleFilterState:
     """Convert initial SMC state particles from `(N, F, ...)` to `(F, N, ...)`.
 
@@ -69,7 +69,6 @@ def factorialize_init_state(
 
     Args:
         init_state: Output from particle filter `init_prepare`
-        model_inputs: The model inputs at the first time point - unused.
     """
     particles = tree.map(lambda x: jnp.moveaxis(x, 0, 1), init_state.particles)
     n_factors = tree.leaves(particles)[0].shape[0]
diff --git a/cuthbert/factorial/types.py b/cuthbert/factorial/types.py
index ea67b499..b2f31ce9 100644
--- a/cuthbert/factorial/types.py
+++ b/cuthbert/factorial/types.py
@@ -143,14 +143,11 @@ def __call__(
 class FactorializeInitState(Protocol):
     """Protocol for factorial post-processing of `init_prepare`."""
 
-    def __call__(
-        self, init_state: ArrayTreeLike, model_inputs: ArrayTreeLike
-    ) -> ArrayTree:
+    def __call__(self, init_state: ArrayTreeLike) -> ArrayTree:
         """Any processing of the output of `init_prepare` for factorial inference.
 
         Args:
             init_state: Output from base inference method's `init_prepare`
-            model_inputs: The model inputs at the first time point.
         """
         ...
 
@@ -182,9 +179,7 @@ class Factorializer(NamedTuple):
     join: Join
     marginalize: Marginalize
     insert: Insert
-    factorialize_init_state: FactorializeInitState = lambda init_state, model_inputs: (
-        init_state
-    )
+    factorialize_init_state: FactorializeInitState = lambda init_state: init_state
 
     def extract_and_join(
         self, factorial_state: ArrayTreeLike, model_inputs: ArrayTreeLike
diff --git a/cuthbert/factorial/utils.py b/cuthbert/factorial/utils.py
index 0e40f023..ad59b403 100644
--- a/cuthbert/factorial/utils.py
+++ b/cuthbert/factorial/utils.py
@@ -4,6 +4,7 @@
 from jax import tree, vmap
 
 from cuthbert.factorial.types import Extract
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayLike, ArrayTree, ArrayTreeLike
 
 
@@ -30,7 +31,9 @@ def serial_to_factorial(
             maximum index, preserving the behaviour of earlier versions.
         init_factorial_tree: Optional initial factorial tree to use, as the first
             element of each returned tree.
-            Leaves with shape (F, ...)
+            Leaves with shape (F, ...). Subtrees that are None in the initial
+            tree (such as model_inputs) are padded with dummy values matching
+            the history.
 
     Returns:
         A single tree when a single index is selected; otherwise, a list of trees in the
@@ -72,12 +75,23 @@ def extract_single_factor(factorial_index: ArrayLike) -> ArrayTree:
             return factor_states
 
         initial_state = extract(init_factorial_tree, factorial_index)
+
+        def prepend_initial(initial_leaf, history_subtree):
+            if initial_leaf is None:
+                # Subtrees absent from the initial tree (such as model_inputs)
+                # are padded with dummy values matching the history.
+                initial_leaf = dummy_leading_element(history_subtree)
+            return tree.map(
+                lambda initial, history: jnp.concatenate([initial[None], history]),
+                initial_leaf,
+                history_subtree,
+            )
+
         return tree.map(
-            lambda initial_leaf, history_leaf: jnp.concatenate(
-                [initial_leaf[None], history_leaf]
-            ),
+            prepend_initial,
             initial_state,
             factor_states,
+            is_leaf=lambda x: x is None,
         )
 
     factorial_trees = [
diff --git a/cuthbert/filtering.py b/cuthbert/filtering.py
index 4c6008a5..72728531 100644
--- a/cuthbert/filtering.py
+++ b/cuthbert/filtering.py
@@ -7,7 +7,7 @@
 from jax.lax import associative_scan, scan
 
 from cuthbert.inference import Filter
-from cuthbert.utils import dummy_tree_like
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -54,8 +54,7 @@ def filter(
     else:
         prepare_keys = random.split(key, T)
 
-    sample_model_inputs = tree.map(lambda x: x[0], model_inputs)
-    init_state = init_state._replace(model_inputs=dummy_tree_like(sample_model_inputs))
+    init_state = init_state._replace(model_inputs=dummy_leading_element(model_inputs))
 
     if parallel:
         other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
diff --git a/cuthbert/gaussian/kalman.py b/cuthbert/gaussian/kalman.py
index e846aa56..3a09944d 100644
--- a/cuthbert/gaussian/kalman.py
+++ b/cuthbert/gaussian/kalman.py
@@ -13,7 +13,7 @@
 from cuthbert.inference import Filter, Smoother
 from cuthbert.utils import dummy_tree_like
 from cuthbertlib.kalman import filtering, smoothing
-from cuthbertlib.types import Array, ArrayTree, ArrayTreeLike, KeyArray
+from cuthbertlib.types import Array, ArrayLike, ArrayTree, ArrayTreeLike, KeyArray
 
 
 class KalmanFilterState(NamedTuple):
@@ -58,8 +58,8 @@ def chol_cov(self) -> Array:
 
 
 def build_filter(
-    m0: Array,
-    chol_P0: Array,
+    m0: ArrayLike,
+    chol_P0: ArrayLike,
     get_dynamics_params: GetDynamicsParams,
     get_observation_params: GetObservationParams,
 ) -> Filter:
@@ -127,8 +127,8 @@ def build_smoother(
 
 
 def init_prepare(
-    m0: Array,
-    chol_P0: Array,
+    m0: ArrayLike,
+    chol_P0: ArrayLike,
     key: KeyArray | None = None,
 ) -> KalmanFilterState:
     """Prepare the initial state for the Kalman filter.
@@ -142,6 +142,7 @@ def init_prepare(
     Returns:
         State for the Kalman filter.
     """
+    m0, chol_P0 = jnp.asarray(m0), jnp.asarray(chol_P0)
     elem = filtering.FilterScanElement(
         A=jnp.zeros_like(chol_P0),
         b=m0,
diff --git a/cuthbert/gaussian/moments/associative_filter.py b/cuthbert/gaussian/moments/associative_filter.py
index 3b474d52..b074e327 100644
--- a/cuthbert/gaussian/moments/associative_filter.py
+++ b/cuthbert/gaussian/moments/associative_filter.py
@@ -1,27 +1,26 @@
 """Implements the associative linearized moments Kalman filter."""
 
-from jax import eval_shape, tree
+from jax import ShapeDtypeStruct, tree
 from jax import numpy as jnp
 
-from cuthbert.gaussian.kalman import GetInitParams
 from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
 from cuthbert.gaussian.types import LinearizedKalmanFilterState
 from cuthbert.utils import dummy_tree_like
 from cuthbertlib.kalman import filtering
 from cuthbertlib.linearize import linearize_moments
-from cuthbertlib.types import ArrayTreeLike, KeyArray
+from cuthbertlib.types import ArrayLike, ArrayTreeLike, KeyArray
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
-    get_init_params: GetInitParams,
+    m0: ArrayLike,
+    chol_P0: ArrayLike,
     key: KeyArray | None = None,
 ) -> LinearizedKalmanFilterState:
     """Prepare the initial state for the linearized moments Kalman filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_params: Function to get m0, chol_P0 from model inputs.
+        m0: The mean vector of the initial state.
+        chol_P0: The generalized Cholesky factor of the initial covariance.
         key: JAX random key - not used.
 
     Returns:
@@ -29,8 +28,7 @@ def init_prepare(
             Contains mean, chol_cov (generalised Cholesky factor of covariance)
             and log_normalizing_constant.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    m0, chol_P0 = get_init_params(model_inputs)
+    m0, chol_P0 = jnp.asarray(m0), jnp.asarray(chol_P0)
 
     prior_state = LinearizedKalmanFilterState(
         elem=filtering.FilterScanElement(
@@ -41,7 +39,7 @@ def init_prepare(
             Z=jnp.zeros_like(chol_P0),
             ell=jnp.array(0.0),
         ),
-        model_inputs=model_inputs,
+        model_inputs=None,
         mean_prev=dummy_tree_like(m0),
     )
 
@@ -50,7 +48,7 @@ def init_prepare(
 
 def filter_prepare(
     model_inputs: ArrayTreeLike,
-    get_init_params: GetInitParams,
+    m0: ArrayLike,
     get_dynamics_params: GetDynamicsMoments,
     get_observation_params: GetObservationMoments,
     key: KeyArray | None = None,
@@ -64,8 +62,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        get_init_params: Function to get m0, chol_P0 from model inputs.
-            Only used to infer shape of mean and chol_cov.
+        m0: The mean vector of the initial state.
+            Only used to infer the state shape.
         get_dynamics_params: Function to get dynamics conditional mean and
             (generalised) Cholesky covariance from linearization point and model inputs.
             `associative_scan` only supported when `state` is ignored.
@@ -79,10 +77,9 @@ def filter_prepare(
         Prepared state for linearized moments Kalman filter.
     """
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    dummy_mean_struct = eval_shape(lambda mi: get_init_params(mi)[0], model_inputs)
-    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_mean = dummy_tree_like(jnp.asarray(m0))
     dummy_chol_cov = dummy_tree_like(
-        jnp.empty(dummy_mean.shape + dummy_mean.shape[-1:], dtype=dummy_mean.dtype)
+        ShapeDtypeStruct(dummy_mean.shape + dummy_mean.shape[-1:], dummy_mean.dtype)
     )
 
     dummy_state = LinearizedKalmanFilterState(
diff --git a/cuthbert/gaussian/moments/filter.py b/cuthbert/gaussian/moments/filter.py
index 344f101a..8897c631 100644
--- a/cuthbert/gaussian/moments/filter.py
+++ b/cuthbert/gaussian/moments/filter.py
@@ -20,12 +20,13 @@
 
 from cuthbert.gaussian.moments import associative_filter, non_associative_filter
 from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
-from cuthbert.gaussian.types import GetInitParams
 from cuthbert.inference import Filter
+from cuthbertlib.types import ArrayLike
 
 
 def build_filter(
-    get_init_params: GetInitParams,
+    m0: ArrayLike,
+    chol_P0: ArrayLike,
     get_dynamics_params: GetDynamicsMoments,
     get_observation_params: GetObservationMoments,
     associative: bool = False,
@@ -41,7 +42,8 @@ def build_filter(
     observation parameters.
 
     Args:
-        get_init_params: Function to get m0, chol_P0 from model inputs.
+        m0: The mean vector of the initial state.
+        chol_P0: The generalized Cholesky factor of the initial covariance.
         get_dynamics_params: Function to get dynamics conditional mean and
             (generalised) Cholesky covariance from linearization point and model inputs.
             and linearization points (for the previous and current time points)
@@ -62,11 +64,11 @@ def build_filter(
     if associative:
         return Filter(
             init_prepare=partial(
-                associative_filter.init_prepare, get_init_params=get_init_params
+                associative_filter.init_prepare, m0=m0, chol_P0=chol_P0
             ),
             filter_prepare=partial(
                 associative_filter.filter_prepare,
-                get_init_params=get_init_params,
+                m0=m0,
                 get_dynamics_params=get_dynamics_params,
                 get_observation_params=get_observation_params,
             ),
@@ -76,11 +78,9 @@ def build_filter(
     else:
         return Filter(
             init_prepare=partial(
-                non_associative_filter.init_prepare, get_init_params=get_init_params
-            ),
-            filter_prepare=partial(
-                non_associative_filter.filter_prepare, get_init_params=get_init_params
+                non_associative_filter.init_prepare, m0=m0, chol_P0=chol_P0
             ),
+            filter_prepare=partial(non_associative_filter.filter_prepare, m0=m0),
             filter_combine=partial(
                 non_associative_filter.filter_combine,
                 get_dynamics_params=get_dynamics_params,
diff --git a/cuthbert/gaussian/moments/non_associative_filter.py b/cuthbert/gaussian/moments/non_associative_filter.py
index efe11afe..a4ce2fec 100644
--- a/cuthbert/gaussian/moments/non_associative_filter.py
+++ b/cuthbert/gaussian/moments/non_associative_filter.py
@@ -1,27 +1,27 @@
 """Implements the non-associative linearized moments Kalman filter."""
 
-from jax import eval_shape, tree
+from jax import ShapeDtypeStruct, tree
 from jax import numpy as jnp
 
 from cuthbert.gaussian.moments.types import GetDynamicsMoments, GetObservationMoments
-from cuthbert.gaussian.types import GetInitParams, LinearizedKalmanFilterState
+from cuthbert.gaussian.types import LinearizedKalmanFilterState
 from cuthbert.gaussian.utils import linearized_kalman_filter_state_dummy_elem
 from cuthbert.utils import dummy_tree_like
 from cuthbertlib.kalman import filtering
 from cuthbertlib.linearize import linearize_moments
-from cuthbertlib.types import ArrayTreeLike, KeyArray
+from cuthbertlib.types import ArrayLike, ArrayTreeLike, KeyArray
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
-    get_init_params: GetInitParams,
+    m0: ArrayLike,
+    chol_P0: ArrayLike,
     key: KeyArray | None = None,
 ) -> LinearizedKalmanFilterState:
     """Prepare the initial state for the linearized moments Kalman filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_params: Function to get m0, chol_P0 from model inputs.
+        m0: The mean vector of the initial state.
+        chol_P0: The generalized Cholesky factor of the initial covariance.
         key: JAX random key - not used.
 
     Returns:
@@ -29,8 +29,7 @@ def init_prepare(
             Contains mean, chol_cov (generalised Cholesky factor of covariance)
             and log_normalizing_constant.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    m0, chol_P0 = get_init_params(model_inputs)
+    m0, chol_P0 = jnp.asarray(m0), jnp.asarray(chol_P0)
 
     prior_state = LinearizedKalmanFilterState(
         elem=filtering.FilterScanElement(
@@ -41,7 +40,7 @@ def init_prepare(
             Z=jnp.zeros_like(chol_P0),
             ell=jnp.array(0.0),
         ),
-        model_inputs=model_inputs,
+        model_inputs=None,
         mean_prev=dummy_tree_like(m0),
     )
 
@@ -50,7 +49,7 @@ def init_prepare(
 
 def filter_prepare(
     model_inputs: ArrayTreeLike,
-    get_init_params: GetInitParams,
+    m0: ArrayLike,
     key: KeyArray | None = None,
 ) -> LinearizedKalmanFilterState:
     """Prepare a state for a linearized moments Kalman filter step.
@@ -59,18 +58,17 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        get_init_params: Function to get m0, chol_P0 from model inputs.
-            Only used to infer shape of mean and chol_cov.
+        m0: The mean vector of the initial state.
+            Only used to infer the state shape.
         key: JAX random key - not used.
 
     Returns:
         Prepared state for linearized moments Kalman filter.
     """
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    dummy_mean_struct = eval_shape(lambda mi: get_init_params(mi)[0], model_inputs)
-    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_mean = dummy_tree_like(jnp.asarray(m0))
     dummy_chol_cov = dummy_tree_like(
-        jnp.empty(dummy_mean.shape + dummy_mean.shape[-1:], dtype=dummy_mean.dtype)
+        ShapeDtypeStruct(dummy_mean.shape + dummy_mean.shape[-1:], dummy_mean.dtype)
     )
 
     return linearized_kalman_filter_state_dummy_elem(
diff --git a/cuthbert/gaussian/taylor/__init__.py b/cuthbert/gaussian/taylor/__init__.py
index be14a3a2..d03d1637 100644
--- a/cuthbert/gaussian/taylor/__init__.py
+++ b/cuthbert/gaussian/taylor/__init__.py
@@ -7,7 +7,6 @@
 from cuthbert.gaussian.taylor.smoother import build_smoother
 from cuthbert.gaussian.taylor.types import (
     GetDynamicsLogDensity,
-    GetInitLogDensity,
     GetObservationFunc,
     LogPotential,
 )
diff --git a/cuthbert/gaussian/taylor/associative_filter.py b/cuthbert/gaussian/taylor/associative_filter.py
index 43c33ff6..7a72a27d 100644
--- a/cuthbert/gaussian/taylor/associative_filter.py
+++ b/cuthbert/gaussian/taylor/associative_filter.py
@@ -1,24 +1,23 @@
 """Implements the associative linearized Taylor Kalman filter."""
 
-from jax import eval_shape, tree
+from jax import ShapeDtypeStruct, tree
 from jax import numpy as jnp
 
 from cuthbert.gaussian.taylor import non_associative_filter
 from cuthbert.gaussian.taylor.types import (
     GetDynamicsLogDensity,
-    GetInitLogDensity,
     GetObservationFunc,
 )
 from cuthbert.gaussian.types import LinearizedKalmanFilterState
 from cuthbert.utils import dummy_tree_like
 from cuthbertlib.kalman import filtering
 from cuthbertlib.linearize import linearize_log_density
-from cuthbertlib.types import ArrayTreeLike, KeyArray
+from cuthbertlib.types import ArrayLike, ArrayTreeLike, KeyArray, LogDensity
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
-    get_init_log_density: GetInitLogDensity,
+    init_log_density: LogDensity,
+    init_linearization_point: ArrayLike,
     rtol: float | None = None,
     ignore_nan_dims: bool = False,
     key: KeyArray | None = None,
@@ -26,9 +25,9 @@ def init_prepare(
     """Prepare the initial state for the linearized Taylor Kalman filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_log_density: Function that returns log density log p(x_0)
-            and linearization point.
+        init_log_density: Initial log density log p(x_0). For factorial models,
+            this is the sum of the log densities for all factors.
+        init_linearization_point: Linearization point for the initial log density.
         rtol: The relative tolerance for the singular values of precision matrices
             when passed to `symmetric_inv_sqrt` during linearization.
             Cutoff for small singular values; singular values smaller than
@@ -46,7 +45,7 @@ def init_prepare(
             and log_normalizing_constant.
     """
     non_assoc_init_state = non_associative_filter.init_prepare(
-        model_inputs, get_init_log_density, rtol, ignore_nan_dims, key
+        init_log_density, init_linearization_point, rtol, ignore_nan_dims, key
     )
     m0 = non_assoc_init_state.mean
     chol_P0 = non_assoc_init_state.chol_cov
@@ -60,7 +59,7 @@ def init_prepare(
             Z=jnp.zeros_like(chol_P0),
             ell=jnp.array(0.0),
         ),
-        model_inputs=model_inputs,
+        model_inputs=None,
         mean_prev=dummy_tree_like(m0),
     )
     return prior_state
@@ -68,7 +67,7 @@ def init_prepare(
 
 def filter_prepare(
     model_inputs: ArrayTreeLike,
-    get_init_log_density: GetInitLogDensity,
+    array_to_infer_shape: ArrayLike,
     get_dynamics_log_density: GetDynamicsLogDensity,
     get_observation_func: GetObservationFunc,
     rtol: float | None = None,
@@ -82,8 +81,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        get_init_log_density: Function that returns log density log p(x_0)
-            and linearization point. Only used to infer shape of mean and chol_cov.
+        array_to_infer_shape: Array with shape matching state or mean, used to infer the
+            state shape.
         get_dynamics_log_density: Function to get dynamics log density log p(x_t+1 | x_t)
             and linearization points (for the previous and current time points)
             `associative_scan` only supported when `state` is ignored.
@@ -106,10 +105,9 @@ def filter_prepare(
         Prepared state for linearized Taylor Kalman filter.
     """
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    dummy_mean_struct = eval_shape(lambda mi: get_init_log_density(mi)[1], model_inputs)
-    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_mean = dummy_tree_like(jnp.asarray(array_to_infer_shape))
     dummy_chol_cov = dummy_tree_like(
-        jnp.empty(dummy_mean.shape + dummy_mean.shape[-1:], dtype=dummy_mean.dtype)
+        ShapeDtypeStruct(dummy_mean.shape + dummy_mean.shape[-1:], dummy_mean.dtype)
     )
 
     dummy_state = LinearizedKalmanFilterState(
diff --git a/cuthbert/gaussian/taylor/filter.py b/cuthbert/gaussian/taylor/filter.py
index f900edb0..4b0f11a0 100644
--- a/cuthbert/gaussian/taylor/filter.py
+++ b/cuthbert/gaussian/taylor/filter.py
@@ -32,14 +32,15 @@
 from cuthbert.gaussian.taylor import associative_filter, non_associative_filter
 from cuthbert.gaussian.taylor.types import (
     GetDynamicsLogDensity,
-    GetInitLogDensity,
     GetObservationFunc,
 )
 from cuthbert.inference import Filter
+from cuthbertlib.types import ArrayLike, LogDensity
 
 
 def build_filter(
-    get_init_log_density: GetInitLogDensity,
+    init_log_density: LogDensity,
+    init_linearization_point: ArrayLike,
     get_dynamics_log_density: GetDynamicsLogDensity,
     get_observation_func: GetObservationFunc,
     associative: bool = False,
@@ -57,9 +58,9 @@ def build_filter(
     observation parameters.
 
     Args:
-        get_init_log_density: Function to get log density log p(x_0)
-            and linearization point.
-            Only takes `model_inputs` as input.
+        init_log_density: Initial log density log p(x_0). For factorial models,
+            this is the sum of the log densities for all factors.
+        init_linearization_point: Linearization point for the initial log density.
         get_dynamics_log_density: Function to get dynamics log density log p(x_t+1 | x_t)
             and linearization points (for the previous and current time points)
             If `associative` is True, the `state` argument should be ignored.
@@ -89,13 +90,14 @@ def build_filter(
         return Filter(
             init_prepare=partial(
                 associative_filter.init_prepare,
-                get_init_log_density=get_init_log_density,
+                init_log_density=init_log_density,
+                init_linearization_point=init_linearization_point,
                 rtol=rtol,
                 ignore_nan_dims=ignore_nan_dims,
             ),
             filter_prepare=partial(
                 associative_filter.filter_prepare,
-                get_init_log_density=get_init_log_density,
+                array_to_infer_shape=init_linearization_point,
                 get_dynamics_log_density=get_dynamics_log_density,
                 get_observation_func=get_observation_func,
                 rtol=rtol,
@@ -108,13 +110,14 @@ def build_filter(
         return Filter(
             init_prepare=partial(
                 non_associative_filter.init_prepare,
-                get_init_log_density=get_init_log_density,
+                init_log_density=init_log_density,
+                init_linearization_point=init_linearization_point,
                 rtol=rtol,
                 ignore_nan_dims=ignore_nan_dims,
             ),
             filter_prepare=partial(
                 non_associative_filter.filter_prepare,
-                get_init_log_density=get_init_log_density,
+                array_to_infer_shape=init_linearization_point,
             ),
             filter_combine=partial(
                 non_associative_filter.filter_combine,
diff --git a/cuthbert/gaussian/taylor/non_associative_filter.py b/cuthbert/gaussian/taylor/non_associative_filter.py
index 0ea33da7..238b9f84 100644
--- a/cuthbert/gaussian/taylor/non_associative_filter.py
+++ b/cuthbert/gaussian/taylor/non_associative_filter.py
@@ -1,11 +1,10 @@
 """Implements the non-associative linearized Taylor Kalman filter."""
 
-from jax import eval_shape, tree
+from jax import ShapeDtypeStruct, tree
 from jax import numpy as jnp
 
 from cuthbert.gaussian.taylor.types import (
     GetDynamicsLogDensity,
-    GetInitLogDensity,
     GetObservationFunc,
     LogConditionalDensity,
     LogPotential,
@@ -16,7 +15,7 @@
 from cuthbertlib.kalman import filtering
 from cuthbertlib.linalg import block_marginal_sqrt_cov
 from cuthbertlib.linearize import linearize_log_density, linearize_taylor
-from cuthbertlib.types import Array, ArrayTreeLike, KeyArray
+from cuthbertlib.types import Array, ArrayLike, ArrayTreeLike, KeyArray, LogDensity
 
 
 def process_observation(
@@ -59,8 +58,8 @@ def process_observation(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
-    get_init_log_density: GetInitLogDensity,
+    init_log_density: LogDensity,
+    init_linearization_point: ArrayLike,
     rtol: float | None = None,
     ignore_nan_dims: bool = False,
     key: KeyArray | None = None,
@@ -68,9 +67,9 @@ def init_prepare(
     """Prepare the initial state for the linearized Taylor Kalman filter.
 
     Args:
-        model_inputs: Model inputs.
-        get_init_log_density: Function that returns log density log p(x_0)
-            and linearization point.
+        init_log_density: Initial log density log p(x_0). For factorial models,
+            this is the sum of the log densities for all factors.
+        init_linearization_point: Linearization point for the initial log density.
         rtol: The relative tolerance for the singular values of precision matrices
             when passed to `symmetric_inv_sqrt` during linearization.
             Cutoff for small singular values; singular values smaller than
@@ -87,8 +86,7 @@ def init_prepare(
             Contains mean, chol_cov (generalised Cholesky factor of covariance)
             and log_normalizing_constant.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    init_log_density, linearization_point = get_init_log_density(model_inputs)
+    linearization_point = jnp.asarray(init_linearization_point)
 
     # Handle factorial axis if present
     linearization_shape = linearization_point.shape
@@ -113,7 +111,7 @@ def flat_init_log_density(_, flat_x):
         mean=m0,
         chol_cov=chol_P0,
         log_normalizing_constant=jnp.array(0.0),
-        model_inputs=model_inputs,
+        model_inputs=None,
         mean_prev=dummy_tree_like(m0),
     )
     return prior_state
@@ -121,7 +119,7 @@ def flat_init_log_density(_, flat_x):
 
 def filter_prepare(
     model_inputs: ArrayTreeLike,
-    get_init_log_density: GetInitLogDensity,
+    array_to_infer_shape: ArrayLike,
     key: KeyArray | None = None,
 ) -> LinearizedKalmanFilterState:
     """Prepare a state for a linearized Taylor Kalman filter step.
@@ -130,18 +128,17 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        get_init_log_density: Function that returns log density log p(x_0)
-            and linearization point. Only used to infer shape of mean and chol_cov.
+        array_to_infer_shape: Array with shape matching state or mean, used to infer the
+            state shape.
         key: JAX random key - not used.
 
     Returns:
         Prepared state for linearized Taylor Kalman filter.
     """
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
-    dummy_mean_struct = eval_shape(lambda mi: get_init_log_density(mi)[1], model_inputs)
-    dummy_mean = dummy_tree_like(dummy_mean_struct)
+    dummy_mean = dummy_tree_like(jnp.asarray(array_to_infer_shape))
     dummy_chol_cov = dummy_tree_like(
-        jnp.empty(dummy_mean.shape + dummy_mean.shape[-1:], dtype=dummy_mean.dtype)
+        ShapeDtypeStruct(dummy_mean.shape + dummy_mean.shape[-1:], dummy_mean.dtype)
     )
 
     return linearized_kalman_filter_state_dummy_elem(
diff --git a/cuthbert/gaussian/taylor/types.py b/cuthbert/gaussian/taylor/types.py
index fd1f9def..b722ddae 100644
--- a/cuthbert/gaussian/taylor/types.py
+++ b/cuthbert/gaussian/taylor/types.py
@@ -15,24 +15,6 @@
 LogPotential: TypeAlias = LogDensity
 
 
-class GetInitLogDensity(Protocol):
-    """Protocol for extracting the initial specifications."""
-
-    def __call__(self, model_inputs: ArrayTreeLike) -> tuple[LogDensity, Array]:
-        """Get the initial log density and initial linearization point.
-
-        If the state-space model is a factorial model, this should return the
-        sum of the log densities for each factor.
-
-        Args:
-            model_inputs: Model inputs.
-
-        Returns:
-            Tuple with initial log density and initial linearization point.
-        """
-        ...
-
-
 class GetDynamicsLogDensity(Protocol):
     """Protocol for extracting the dynamics specifications."""
 
diff --git a/cuthbert/inference.py b/cuthbert/inference.py
index 45c90506..e30dca18 100644
--- a/cuthbert/inference.py
+++ b/cuthbert/inference.py
@@ -11,7 +11,10 @@ class InitPrepare(Protocol):
     def __call__(self, *, key: KeyArray | None = None) -> ArrayTree:
         """Prepare the initial state for the inference.
 
-        The state at the first time point, prior to any observations.
+        The state at the first time point, prior to any observations. Initial
+        parameters are supplied when building the filter. The returned state has
+        `model_inputs=None`; offline filtering pads this field to match the step
+        input tree.
 
         Args:
             key: The key for the random number generator.
diff --git a/cuthbert/smc/marginal_particle_filter.py b/cuthbert/smc/marginal_particle_filter.py
index c3336651..649c8f71 100644
--- a/cuthbert/smc/marginal_particle_filter.py
+++ b/cuthbert/smc/marginal_particle_filter.py
@@ -38,7 +38,8 @@ def build_filter(
     r"""Builds a marginal particle filter object.
 
     Args:
-        init_sample: Function to sample from the initial distribution $M_0(x_0)$.
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$.
         propagate_sample: Function to sample from the Markov kernel $M_t(x_t \mid x_{t-1})$.
         log_potential: Function to compute the log potential $\log G_t(x_{t-1}, x_t)$.
         n_filter_particles: Number of particles for the filter.
@@ -72,7 +73,6 @@ def build_filter(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
     init_sample: InitSample,
     n_filter_particles: int,
     key: KeyArray | None = None,
@@ -80,8 +80,8 @@ def init_prepare(
     """Prepare the initial state for the marginal particle filter.
 
     Args:
-        model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution M_0(x_0).
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$.
         n_filter_particles: Number of particles to sample.
         key: JAX random key.
 
@@ -91,13 +91,12 @@ def init_prepare(
     Raises:
         ValueError: If `key` is None.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Sample
     keys = random.split(key, n_filter_particles)
-    particles = jax.vmap(init_sample, (0, None))(keys, model_inputs)
+    particles = jax.vmap(init_sample)(keys)
 
     # Weight
     log_weights = jnp.zeros(n_filter_particles)
@@ -111,7 +110,7 @@ def init_prepare(
         key=key,
         particles=particles,
         log_weights=log_weights,
-        model_inputs=model_inputs,
+        model_inputs=None,
         log_normalizing_constant=log_normalizing_constant,
     )
 
@@ -126,8 +125,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution M_0(x_0).
-            Only used to infer particle shapes.
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$. Only used to infer particle shapes.
         n_filter_particles: Number of particles for the filter.
         key: JAX random key.
 
@@ -140,11 +139,13 @@ def filter_prepare(
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
-    dummy_particle = jax.eval_shape(init_sample, key, model_inputs)
-    particles = tree.map(
-        lambda x: jnp.empty((n_filter_particles,) + x.shape), dummy_particle
+    dummy_particle = jax.eval_shape(init_sample, key)
+    particles = dummy_tree_like(
+        tree.map(
+            lambda x: jax.ShapeDtypeStruct((n_filter_particles,) + x.shape, x.dtype),
+            dummy_particle,
+        )
     )
-    particles = dummy_tree_like(particles)
     return MarginalParticleFilterState(
         key=key,
         particles=particles,
diff --git a/cuthbert/smc/particle_filter.py b/cuthbert/smc/particle_filter.py
index d445af44..8feac4ae 100644
--- a/cuthbert/smc/particle_filter.py
+++ b/cuthbert/smc/particle_filter.py
@@ -43,7 +43,8 @@ def build_filter(
     r"""Builds a particle filter object.
 
     Args:
-        init_sample: Function to sample from the initial distribution $M_0(x_0)$.
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$.
         propagate_sample: Function to sample from the Markov kernel $M_t(x_t \mid x_{t-1})$.
         log_potential: Function to compute the log potential $\log G_t(x_{t-1}, x_t)$.
         n_filter_particles: Number of particles for the filter.
@@ -77,7 +78,6 @@ def build_filter(
 
 
 def init_prepare(
-    model_inputs: ArrayTreeLike,
     init_sample: InitSample,
     n_filter_particles: int,
     key: KeyArray | None = None,
@@ -85,8 +85,8 @@ def init_prepare(
     """Prepare the initial state for the particle filter.
 
     Args:
-        model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution M_0(x_0).
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$.
         n_filter_particles: Number of particles to sample.
         key: JAX random key.
 
@@ -96,13 +96,12 @@ def init_prepare(
     Raises:
         ValueError: If `key` is None.
     """
-    model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
 
     # Sample
     keys = random.split(key, n_filter_particles)
-    particles = jax.vmap(init_sample, (0, None))(keys, model_inputs)
+    particles = jax.vmap(init_sample)(keys)
 
     # Weight
     log_weights = jnp.zeros(n_filter_particles)
@@ -117,7 +116,7 @@ def init_prepare(
         particles=particles,
         log_weights=log_weights,
         ancestor_indices=jnp.arange(n_filter_particles),
-        model_inputs=model_inputs,
+        model_inputs=None,
         log_normalizing_constant=log_normalizing_constant,
     )
 
@@ -132,8 +131,8 @@ def filter_prepare(
 
     Args:
         model_inputs: Model inputs.
-        init_sample: Function to sample from the initial distribution M_0(x_0).
-            Only used to infer particle shapes.
+        init_sample: Function of a JAX random key only, generates a single sample
+            from $M_0(x_0)$. Only used to infer particle shapes.
         n_filter_particles: Number of particles for the filter.
         key: JAX random key.
 
@@ -146,11 +145,13 @@ def filter_prepare(
     model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs)
     if key is None:
         raise ValueError("A JAX PRNG key must be provided.")
-    dummy_particle = jax.eval_shape(init_sample, key, model_inputs)
-    particles = tree.map(
-        lambda x: jnp.empty((n_filter_particles,) + x.shape), dummy_particle
+    dummy_particle = jax.eval_shape(init_sample, key)
+    particles = dummy_tree_like(
+        tree.map(
+            lambda x: jax.ShapeDtypeStruct((n_filter_particles,) + x.shape, x.dtype),
+            dummy_particle,
+        )
     )
-    particles = dummy_tree_like(particles)
     return ParticleFilterState(
         key=key,
         particles=particles,
diff --git a/cuthbert/smc/types.py b/cuthbert/smc/types.py
index ee7af691..8c25c2a2 100644
--- a/cuthbert/smc/types.py
+++ b/cuthbert/smc/types.py
@@ -13,15 +13,14 @@
 class InitSample(Protocol):
     """Protocol for  sampling from the initial distribution $M_0(x_0)$."""
 
-    def __call__(self, key: KeyArray, model_inputs: ArrayTreeLike) -> ArrayTree:
+    def __call__(self, key: KeyArray) -> ArrayTree:
         """Samples from the initial distribution $M_0(x_0)$.
 
         Args:
             key: JAX PRNG key.
-            model_inputs: Model inputs.
 
         Returns:
-            A sample $x_0$.
+            A single sample $x_0$.
         """
         ...
 
diff --git a/cuthbert/smoothing.py b/cuthbert/smoothing.py
index b70a688e..9c486723 100644
--- a/cuthbert/smoothing.py
+++ b/cuthbert/smoothing.py
@@ -7,7 +7,7 @@
 from jax.lax import associative_scan, scan
 
 from cuthbert.inference import Smoother
-from cuthbert.utils import dummy_tree_like
+from cuthbert.utils import dummy_leading_element
 from cuthbertlib.types import ArrayTree, ArrayTreeLike, KeyArray
 
 
@@ -23,11 +23,10 @@ def smoother(
     `filter_states` should have leading temporal dimension of length T + 1, where
     T is the number of time steps excluding the initial state.
 
-    Each element of `model_inputs` refers to the transition from t to t+1, except for the
-    first element which refers to the initial state. The initial state `model_inputs`
-    are not used for smoothing. Thus the `model_inputs` used here have length T.
-    By default, `filter_states.model_inputs[1:]` are used (i.e. the `model_inputs`
-    used for the initial state is ignored).
+    Each of the T elements of `model_inputs` describes a transition from t to t+1.
+    Inputs of length T + 1 are also accepted; the first element is ignored.
+    By default, `filter_states.model_inputs[1:]` are used, discarding the dummy
+    inputs stored alongside the initial state.
 
     Args:
         smoother_obj: The smoother inference object.
@@ -54,7 +53,7 @@ def smoother(
 
     # model_inputs for the dynamics distribution from t-1 to t is stored
     # in model_inputs[t] thus we need model_inputs[1:]
-    # model_inputs[0] is only used for init_prepare and not for smoothing.
+    # model_inputs[0] contains dummy values for the initial state.
     # Therefore, we allow model_inputs to be either of length T + 1 or T
     # where if length is T + 1 then we simply discard model_inputs[0]
     model_inputs_length = tree.leaves(model_inputs)[0].shape[0]
@@ -78,7 +77,7 @@ def smoother(
 
     # Final smoother state doesn't need model inputs, so we create a dummy one
     # with the same structure as model_inputs but with all values set to dummy values.
-    dummy_single_model_inputs = dummy_tree_like(tree.map(lambda x: x[0], model_inputs))
+    dummy_single_model_inputs = dummy_leading_element(model_inputs)
 
     final_smoother_state = smoother_obj.convert_filter_to_smoother_state(
         final_filter_state, model_inputs=dummy_single_model_inputs, key=prepare_keys[0]
diff --git a/cuthbert/utils.py b/cuthbert/utils.py
index 857abb25..e4b2f4ac 100644
--- a/cuthbert/utils.py
+++ b/cuthbert/utils.py
@@ -8,10 +8,9 @@
 
 def _dummy_array(leaf: ArrayLike | jax.ShapeDtypeStruct) -> Array:
     """Returns an array of the same shape and dtype filled with dummy values."""
-    if isinstance(leaf, jax.ShapeDtypeStruct):
-        leaf = jnp.empty_like(leaf)
+    if not isinstance(leaf, jax.ShapeDtypeStruct):
+        leaf = jnp.asarray(leaf)
 
-    leaf = jnp.asarray(leaf)
     dtype = leaf.dtype
     shape = leaf.shape
 
@@ -30,3 +29,24 @@ def _dummy_array(leaf: ArrayLike | jax.ShapeDtypeStruct) -> Array:
 def dummy_tree_like(pytree: ArrayTreeLike) -> ArrayTree:
     """Returns a pytree with the same structure filled with dummy values."""
     return jax.tree.map(_dummy_array, pytree)
+
+
+def dummy_leading_element(pytree: ArrayTreeLike) -> ArrayTree:
+    """Returns dummy values shaped like a single element of the leading axis.
+
+    Only the shape and dtype of each leaf are read, so this is safe for leaves
+    whose leading axis has length zero.
+
+    Args:
+        pytree: Pytree whose leaves all have a leading axis.
+
+    Returns:
+        Pytree with the same structure, each leaf filled with dummy values and
+            with the leading axis removed.
+    """
+    return jax.tree.map(
+        lambda x: _dummy_array(
+            jax.ShapeDtypeStruct(jnp.shape(x)[1:], jnp.result_type(x))
+        ),
+        pytree,
+    )
diff --git a/docs/api_cuthbert/gaussian/kalman.md b/docs/api_cuthbert/gaussian/kalman.md
index 6b787f14..0c86a544 100644
--- a/docs/api_cuthbert/gaussian/kalman.md
+++ b/docs/api_cuthbert/gaussian/kalman.md
@@ -9,6 +9,5 @@
 ::: cuthbert.gaussian.types
     options:
       members:
-        - GetInitParams
         - GetDynamicsParams
         - GetObservationParams
diff --git a/docs/examples/diff_pf_resampling.md b/docs/examples/diff_pf_resampling.md
index 92be34b3..51ab48cb 100644
--- a/docs/examples/diff_pf_resampling.md
+++ b/docs/examples/diff_pf_resampling.md
@@ -110,9 +110,6 @@ def kalman_mll(theta, ys, *, m0, chol_P0):
     d = jnp.array([0.0])
     chol_R = jnp.array([[r]])
 
-    def get_init_params(_mi):
-        return m0, chol_P0
-
     def get_dynamics_params(_mi):
         return F, c, chol_Q
 
@@ -121,9 +118,9 @@ def kalman_mll(theta, ys, *, m0, chol_P0):
         y = ys[idx]
         return H, d, chol_R, y
 
-    kf = kalman.build_filter(get_init_params, get_dynamics_params, get_observation_params)
+    kf = kalman.build_filter(m0, chol_P0, get_dynamics_params, get_observation_params)
     model_inputs = jnp.arange(ys.shape[0] + 1)
-    init_state = kf.init_prepare(model_inputs[0])
+    init_state = kf.init_prepare()
     states = run_filter(kf, model_inputs[1:], init_state, parallel=False)
     return states.log_normalizing_constant[-1]
 
@@ -170,7 +167,7 @@ def pf_mll(
     d = jnp.array([0.0])
     chol_R = jnp.array([[r]])
 
-    def init_sample(k, _mi):
+    def init_sample(k):
         return m0 + (chol_P0 @ random.normal(k, (1,)))
 
     def propagate_sample(k, x_prev, _mi):
@@ -198,7 +195,7 @@ def pf_mll(
 
     model_inputs = jnp.arange(ys.shape[0] + 1)
     init_key, filter_key = jax.random.split(key)
-    init_state = filt.init_prepare(model_inputs[0], key=init_key)
+    init_state = filt.init_prepare(key=init_key)
     states = run_filter(
         filt, model_inputs[1:], init_state, parallel=False, key=filter_key
     )
@@ -515,7 +512,7 @@ Finally, we should be sure that the forward pass is not actually being modified
         d = jnp.array([0.0])
         chol_R = jnp.array([[r]])
 
-        def init_sample(k, _mi):
+        def init_sample(k):
             return m0 + (chol_P0 @ random.normal(k, (1,)))
 
         def propagate_sample(k, x_prev, _mi):
@@ -542,7 +539,7 @@ Finally, we should be sure that the forward pass is not actually being modified
 
         model_inputs = jnp.arange(ys.shape[0] + 1)
         init_key, filter_key = jax.random.split(key)
-        init_state = filt.init_prepare(model_inputs[0], key=init_key)
+        init_state = filt.init_prepare(key=init_key)
         states = run_filter(
             filt, model_inputs[1:], init_state, parallel=False, key=filter_key
         )
diff --git a/docs/examples/dynamax_integration.md b/docs/examples/dynamax_integration.md
index 1db24e18..42c3b862 100644
--- a/docs/examples/dynamax_integration.md
+++ b/docs/examples/dynamax_integration.md
@@ -128,10 +128,6 @@ def build_cuthbert_kalman_filter_from_dynamax(lgssm_model, lgssm_params, observa
     d = lgssm_params.emissions.bias
     chol_R = jnp.linalg.cholesky(lgssm_params.emissions.cov)
 
-    def get_init_params(model_inputs):
-        """Return initial state distribution parameters."""
-        return m0, chol_P0
-
     def get_dynamics_params(model_inputs):
         """Return dynamics parameters.
 
@@ -151,7 +147,7 @@ def build_cuthbert_kalman_filter_from_dynamax(lgssm_model, lgssm_params, observa
 
     # Build the Kalman filter
     filter_obj = kalman.build_filter(
-        get_init_params=get_init_params,
+        m0=m0, chol_P0=chol_P0,
         get_dynamics_params=get_dynamics_params,
         get_observation_params=get_observation_params
     )
@@ -171,7 +167,7 @@ Now we can run the `cuthbert` Kalman filter to obtain the filtering distribution
 
 ```{.python #dynamax-run-filter}
 # Run Kalman filtering
-init_state = filter_obj.init_prepare(model_inputs[0])
+init_state = filter_obj.init_prepare()
 filtered_states = cuthbert.filter(filter_obj, model_inputs[1:], init_state)
 
 # Extract filtering results - remove initial time step
diff --git a/docs/examples/enkf_comparison.md b/docs/examples/enkf_comparison.md
index 091b2cbf..bc70c9d5 100644
--- a/docs/examples/enkf_comparison.md
+++ b/docs/examples/enkf_comparison.md
@@ -127,11 +127,10 @@ model_inputs = jnp.arange(num_time_steps + 1)
 We'll use an EKF through the [`taylor` submodule](../api_cuthbert/gaussian/taylor.md). For this, we need to specify an initial log density, dynamics log-density, observable log-density, and linearization points for each.
 
 ```{.python #enkf-comparison-ekf}
-def get_init_log_density(model_inputs):
-    def init_log_density(x):
-        return logpdf(x, m0, chol_P0, nan_support=False)
+def init_log_density(x):
+    return logpdf(x, m0, chol_P0, nan_support=False)
 
-    return init_log_density, m0
+init_linearization_point = m0
 
 
 def get_dynamics_log_density(
@@ -157,7 +156,8 @@ def get_observation_func(
 
 
 ekf = taylor.build_filter(
-    get_init_log_density,
+    init_log_density,
+    init_linearization_point,
     get_dynamics_log_density,
     get_observation_func,
     associative=False,
@@ -167,7 +167,7 @@ jitted_filter = jit(run_filter, static_argnames=["filter_obj"])
 filter_model_inputs = model_inputs[1:]
 
 n_timing = 20
-ekf_init_state = ekf.init_prepare(model_inputs[0])
+ekf_init_state = ekf.init_prepare()
 ekf_states = jitted_filter(ekf, filter_model_inputs, ekf_init_state)  # warm up
 jax.block_until_ready(ekf_states)
 _times = []
@@ -183,11 +183,11 @@ ekf_chol_covs = ekf_states.chol_cov
 
 ## EnKF
 
-The EnKF propagates an ensemble of particles through the nonlinear dynamics directly. It then performs a Kalman-style update using empirical covariances of these particles. It does not need to compute a Jacobian of the dynamics, unlike the EKF. We need to specify a function that generates initial samples `(key, model_inputs) -> x_0`, a stochastic dynamics simulator `(x, key) -> x_next`, and observation parameters.
+The EnKF propagates an ensemble of particles through the nonlinear dynamics directly. It then performs a Kalman-style update using empirical covariances of these particles. It does not need to compute a Jacobian of the dynamics, unlike the EKF. We need to specify a function that generates initial samples `key -> x_0`, a stochastic dynamics simulator `(x, key) -> x_next`, and observation parameters.
 
 ```{.python #enkf-comparison-enkf}
 enkf = ensemble_kalman_filter.build_filter(
-    init_sample=lambda key, mi: m0 + chol_P0 @ random.normal(key, m0.shape),
+    init_sample=lambda key: m0 + chol_P0 @ random.normal(key, m0.shape),
     get_dynamics=lambda mi: lambda x, key: lorenz_step(x)
     + chol_Q @ random.normal(key, (x_dim,)),
     get_observations=lambda mi: (lambda x: H @ x + d_obs, chol_R, ys[mi - 1]),
@@ -197,7 +197,7 @@ enkf = ensemble_kalman_filter.build_filter(
 )
 
 key, enkf_init_key, enkf_filter_key = random.split(key, 3)
-enkf_init_state = enkf.init_prepare(model_inputs[0], key=enkf_init_key)
+enkf_init_state = enkf.init_prepare(key=enkf_init_key)
 enkf_states = jitted_filter(
     enkf, filter_model_inputs, enkf_init_state, key=enkf_filter_key
 )  # warm up
@@ -222,7 +222,7 @@ The bootstrap particle filter makes no Gaussian assumption. However, empirically
 adaptive_systematic = adaptive.ess_decorator(systematic.resampling, 0.5)
 
 pf = particle_filter.build_filter(
-    init_sample=lambda key, mi: m0 + chol_P0 @ random.normal(key, (x_dim,)),
+    init_sample=lambda key: m0 + chol_P0 @ random.normal(key, (x_dim,)),
     propagate_sample=lambda key, state, mi: lorenz_step(state)
     + chol_Q @ random.normal(key, (x_dim,)),
     log_potential=lambda s_prev, s, mi: logpdf(
@@ -233,7 +233,7 @@ pf = particle_filter.build_filter(
 )
 
 key, pf_init_key, pf_filter_key = random.split(key, 3)
-pf_init_state = pf.init_prepare(model_inputs[0], key=pf_init_key)
+pf_init_state = pf.init_prepare(key=pf_init_key)
 pf_states = jitted_filter(
     pf, filter_model_inputs, pf_init_state, key=pf_filter_key
 )  # warm up
diff --git a/docs/examples/enkf_localization_l96.md b/docs/examples/enkf_localization_l96.md
index 9e2989fc..0f3ff4e1 100644
--- a/docs/examples/enkf_localization_l96.md
+++ b/docs/examples/enkf_localization_l96.md
@@ -184,7 +184,7 @@ The cross-covariance modifier receives the empirical $C_{xy}$, while the innovat
 We now compare a localized and unlocalized (stochastic) EnKF on the generated data. We build each EnKF as follows:
 
 ```{.python #enkf-localization-l96-filter}
-def init_sample(key, _model_inputs):
+def init_sample(key):
     index = random.randint(key, (), 0, climatology_states.shape[0])
     return climatology_states[index]
 
@@ -220,10 +220,9 @@ def build_enkf(
     )
 
 jitted_filter = jax.jit(run_filter, static_argnames=("filter_obj",))
-initial_model_inputs = jnp.full(observation_dim, jnp.nan)
 
 def apply_filter(filter_obj):
-    initial_state = filter_obj.init_prepare(initial_model_inputs, key=init_key)
+    initial_state = filter_obj.init_prepare(key=init_key)
     states = jitted_filter(
         filter_obj,
         observations,
diff --git a/docs/examples/factorial_football.md b/docs/examples/factorial_football.md
index 61da83ab..14095571 100644
--- a/docs/examples/factorial_football.md
+++ b/docs/examples/factorial_football.md
@@ -19,7 +19,7 @@ from jax.scipy.stats import norm
 
 from cuthbert import factorial, smoother
 from cuthbert.gaussian import taylor
-from cuthbertlib.types import LogConditionalDensity, LogDensity
+from cuthbertlib.types import LogConditionalDensity
 ```
 
 Nothing too surprising there I hope. We'll be using the [`taylor`](api_cuthbert/gaussian/taylor.md)
@@ -240,11 +240,10 @@ tau = 0.05
 epsilon = 0.3
 
 
-def get_init_log_density(model_inputs: MatchData) -> tuple[LogDensity, Array]:
-    def init_log_density(x):
-        return norm.logpdf(x, 0, init_sd).sum()
+def init_log_density(x):
+    return norm.logpdf(x, 0, init_sd).sum()
 
-    return init_log_density, jnp.zeros((num_teams, 1))
+init_linearization_point = jnp.zeros((num_teams, 1))
 
 
 def dynamics_log_density(x_prev, x, time_diff):
@@ -308,7 +307,8 @@ Now that we've defined the model, we can construct the `cuthbert` [filter object
 
 ```{.python #factorial-football-build-filter}
 football_filter = taylor.build_filter(
-    get_init_log_density,
+    init_log_density,
+    init_linearization_point,
     get_dynamics_log_density,
     get_observation_func,
 )
@@ -331,11 +331,10 @@ We'll use `cuthbert.factorial.filter` to easily
 run offline filtering on our data.
 
 ```{.python #factorial-football-run-filter}
-init_match_data = tree.map(lambda x: x[0], match_data)
 filter_match_data = tree.map(lambda x: x[1:], match_data)
 
-init_state = football_filter.init_prepare(init_match_data)
-init_state = factorializer.factorialize_init_state(init_state, init_match_data)
+init_state = football_filter.init_prepare()
+init_state = factorializer.factorialize_init_state(init_state)
 
 local_filter_states, final_factorial_state = factorial.filter(
     football_filter, factorializer, filter_match_data, init_state
@@ -403,7 +402,8 @@ def get_dynamics_log_density_single_team(
 
 
 single_team_filter = taylor.build_filter(
-    get_init_log_density,
+    init_log_density,
+    init_linearization_point,
     get_dynamics_log_density_single_team,
     get_observation_func=lambda state, model_inputs: (
         lambda x: jnp.zeros([]),
diff --git a/docs/examples/kalman_tracking.md b/docs/examples/kalman_tracking.md
index 5e52dc17..f8a34a7e 100644
--- a/docs/examples/kalman_tracking.md
+++ b/docs/examples/kalman_tracking.md
@@ -90,9 +90,8 @@ of time-varying parameters.
 
 #### The Parameter Extraction Pattern
 
-Instead of passing matrices directly, `kalman.build_filter()` takes three functions that extract the necessary parameters given a time index (or other model inputs):
+`kalman.build_filter()` takes the initial mean `m0` and covariance factor `chol_P0` directly, followed by two functions that extract step parameters given a time index (or other model inputs):
 
-- **`get_init_params(model_inputs)`**: Returns initial state mean `m0` and noise Cholesky factor `chol_P0`.
 - **`get_dynamics_params(model_inputs)`**: Returns dynamics matrix `F`, bias term `c`, and process noise Cholesky factor `chol_Q`.
 - **`get_observation_params(model_inputs)`**: Returns observation matrix `H`, bias term `d`, observation noise Cholesky factor `chol_R`, and the actual observation `y`.
 
@@ -115,9 +114,6 @@ This pattern separates the model specification from the filtering algorithm, mak
 ```{.python #kalman-build-filter}
 def build_car_tracking_filter(m0, chol_P0, F, c, chol_Q, H, d, chol_R, ys):
 
-    def get_init_params(model_inputs):
-        return m0, chol_P0
-
     def get_dynamics_params(model_inputs):
         return F, c, chol_Q
 
@@ -125,7 +121,7 @@ def build_car_tracking_filter(m0, chol_P0, F, c, chol_Q, H, d, chol_R, ys):
         return H, d, chol_R, ys[model_inputs - 1]
 
     filter_obj = kalman.build_filter(
-        get_init_params, get_dynamics_params, get_observation_params
+        m0, chol_P0, get_dynamics_params, get_observation_params
     )
     model_inputs = jnp.arange(len(ys) + 1)
     return filter_obj, model_inputs
@@ -142,7 +138,7 @@ function with our constructed filter object and model inputs:
 
 ```{.python #kalman-run-filter}
 # Run the filter
-init_state = filter_obj.init_prepare(model_inputs[0])
+init_state = filter_obj.init_prepare()
 filtered_states = filter(filter_obj, model_inputs[1:], init_state, parallel=True)
 
 # Extract results
@@ -216,7 +212,7 @@ car's position.
     )
 
     # Run filtering - cuthbert handles NaNs automatically
-    init_state_missing = filter_obj_missing.init_prepare(model_inputs_missing[0])
+    init_state_missing = filter_obj_missing.init_prepare()
     filtered_states_missing = filter(
         filter_obj_missing, model_inputs_missing[1:], init_state_missing
     )
diff --git a/docs/examples/online_stoch_vol.md b/docs/examples/online_stoch_vol.md
index d2aa9e63..6a026c25 100644
--- a/docs/examples/online_stoch_vol.md
+++ b/docs/examples/online_stoch_vol.md
@@ -154,7 +154,7 @@ Then we'll convert our mathematical model into one that `cuthbert.smc.particle_f
 can understand.
 
 ```{.python #online-stoch-vol-model-func}
-def init_sample(key: Array, model_inputs: ObservationData) -> Array:
+def init_sample(key: Array) -> Array:
     return init_mean + init_std * random.normal(key, ())
 
 
@@ -212,9 +212,8 @@ data.
 
 ```{.python #online-stoch-vol-particle-filter-run-previous}
 key, init_key, previous_key = random.split(random.key(0), 3)
-init_data = tree.map(lambda x: x[0], previous_data)
 filter_data = tree.map(lambda x: x[1:], previous_data)
-init_state = pf.init_prepare(init_data, key=init_key)
+init_state = pf.init_prepare(key=init_key)
 previous_states = filter(pf, filter_data, init_state, key=previous_key)
 filter_state = tree.map(lambda x: x[-1], previous_states)
 ```
diff --git a/docs/examples/parameter_estimation_em.md b/docs/examples/parameter_estimation_em.md
index 5b2040b5..ffc6a4c4 100644
--- a/docs/examples/parameter_estimation_em.md
+++ b/docs/examples/parameter_estimation_em.md
@@ -139,9 +139,6 @@ Let's build the filter and smoother objects below that correspond to the SSM:
 ```{.python #em-model}
 # Build model objects - this is where the model definition is encapsulated
 def model_factory(params: Params):
-    def get_init_params(model_inputs: int) -> tuple[Array, Array]:
-        return jnp.array([0.0]), params.sigma
-
     def get_dynamics_moments(state, model_inputs: int):
         def dynamics_mean_and_chol_cov_func(x):
             mean_t = jnp.where(model_inputs == 0, x, params.rho * x)
@@ -167,7 +164,8 @@ def model_factory(params: Params):
         )
 
     filter_obj = moments.build_filter(
-        get_init_params,
+        jnp.array([0.0]),
+        params.sigma,
         get_dynamics_moments,
         get_observation_moments,
         associative=False,
@@ -298,7 +296,7 @@ n_epochs = 30
 
 for epoch in range(n_epochs):
     filter_obj, smoother_obj = model_factory(params)
-    init_state = filter_obj.init_prepare(model_inputs[0])
+    init_state = filter_obj.init_prepare()
     filtered_states = filter(filter_obj, model_inputs[1:], init_state)
     log_marginal_likelihood_track.append(filtered_states.log_normalizing_constant[-1])
     smoother_states = smoother(smoother_obj, filtered_states)
diff --git a/docs/examples/temporal_parallelization_kalman.md b/docs/examples/temporal_parallelization_kalman.md
index 556b7772..df7bc0d2 100644
--- a/docs/examples/temporal_parallelization_kalman.md
+++ b/docs/examples/temporal_parallelization_kalman.md
@@ -37,9 +37,6 @@ m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys = generate_lgssm(
 )
 
 
-def get_init_params(model_inputs):
-    return m0, chol_P0
-
 
 def get_dynamics_params(model_inputs):
     t = model_inputs - 1
@@ -52,7 +49,7 @@ def get_observation_params(model_inputs):
 
 
 filter_obj = kalman.build_filter(
-    get_init_params, get_dynamics_params, get_observation_params
+    m0, chol_P0, get_dynamics_params, get_observation_params
 )
 model_inputs = jnp.arange(num_time_steps + 1)
 ```
@@ -67,7 +64,7 @@ sequential and parallel implementations.
 
 ```{.python #parallel-kalman-compiletime}
 jitted_filter = jax.jit(filter, static_argnames=("filter_obj", "parallel"))
-init_state = filter_obj.init_prepare(model_inputs[0])
+init_state = filter_obj.init_prepare()
 filter_model_inputs = model_inputs[1:]
 
 seq_compile_time = timeit.Timer(
diff --git a/tests/cuthbert/discrete/test_discrete.py b/tests/cuthbert/discrete/test_discrete.py
index fad2dbb3..e4b924ac 100644
--- a/tests/cuthbert/discrete/test_discrete.py
+++ b/tests/cuthbert/discrete/test_discrete.py
@@ -100,9 +100,6 @@ def std_forward_backward(init_dist, trans_matrix, log_likelihoods):
 def build_inference_object(init_dist, trans_matrices, log_likelihoods):
     T = log_likelihoods.shape[0]
 
-    def get_init_dist(model_inputs):
-        return init_dist
-
     def get_trans_matrix(model_inputs):
         return trans_matrices[model_inputs - 1]
 
@@ -110,7 +107,7 @@ def get_obs_lls(model_inputs):
         return log_likelihoods[model_inputs - 1]
 
     filter_obj = build_filter(
-        get_init_dist=get_init_dist,
+        init_dist=init_dist,
         get_trans_matrix=get_trans_matrix,
         get_obs_lls=get_obs_lls,
     )
@@ -133,7 +130,7 @@ def test(self, seed, num_states, num_time_steps, parallel):
         )
 
         # Run the filter and smoother
-        init_state = filter_obj.init_prepare(model_inputs[0])
+        init_state = filter_obj.init_prepare()
         filtered_states = self.variant(
             filter, static_argnames=("filter_obj", "parallel")
         )(filter_obj, model_inputs[1:], init_state, parallel=parallel)
@@ -178,12 +175,12 @@ def get_noop_obs_lls(model_inputs):
         return jnp.zeros(num_states)
 
     filter_obj = build_filter(
-        get_init_dist=lambda model_inputs: init_dist,
+        init_dist=init_dist,
         get_trans_matrix=get_noop_trans_matrix,
         get_obs_lls=get_noop_obs_lls,
     )
 
-    init_state = filter_obj.init_prepare(None)
+    init_state = filter_obj.init_prepare()
     prep_state = filter_obj.filter_prepare(None)
     filtered_state = filter_obj.filter_combine(init_state, prep_state)
     chex.assert_trees_all_close(filtered_state, init_state, rtol=1e-10, atol=1e-10)
diff --git a/tests/cuthbert/ensemble_kalman/test_enkf.py b/tests/cuthbert/ensemble_kalman/test_enkf.py
index 85e13ad3..4d8a8899 100644
--- a/tests/cuthbert/ensemble_kalman/test_enkf.py
+++ b/tests/cuthbert/ensemble_kalman/test_enkf.py
@@ -41,7 +41,7 @@ def load_enkf_inference(
 ):
     x_dim = m0.shape[0]
 
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         return m0 + chol_P0 @ random.normal(key, m0.shape)
 
     if noop:
@@ -127,7 +127,7 @@ def construct_chol_innovation_covariance(Y, chol_R, model_inputs):
             construct_chol_innovation_covariance=(construct_chol_innovation_covariance),
         )
         init_key, filter_key = random.split(random.key(seed + 1))
-        init_state = inference.init_prepare(model_inputs[0], key=init_key)
+        init_state = inference.init_prepare(key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
             inference, model_inputs[1:], init_state, parallel=False, key=filter_key
         )
@@ -163,7 +163,7 @@ def test_nonlinear_dynamics(self):
             seed, x_dim, y_dim, num_time_steps
         )
 
-        def init_sample(key, model_inputs):
+        def init_sample(key):
             return m0 + chol_P0 @ random.normal(key, m0.shape)
 
         def dynamics_fn(x, key):
@@ -186,7 +186,7 @@ def get_observations(model_inputs):
         model_inputs = jnp.arange(num_time_steps + 1)
 
         init_key, filter_key = random.split(random.key(seed + 1))
-        init_state = inference.init_prepare(model_inputs[0], key=init_key)
+        init_state = inference.init_prepare(key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
             inference, model_inputs[1:], init_state, parallel=False, key=filter_key
         )
@@ -198,7 +198,7 @@ def get_observations(model_inputs):
 
         # Check autodiff works (differentiate w.r.t. a parameter)
         def log_nc(m0_):
-            def init_sample_(key, model_inputs):
+            def init_sample_(key):
                 return m0_ + chol_P0 @ random.normal(key, m0_.shape)
 
             inference_ = ensemble_kalman_filter.build_filter(
@@ -208,7 +208,7 @@ def init_sample_(key, model_inputs):
                 n_particles=1_000,
             )
             init_key, filter_key = random.split(random.key(seed + 1))
-            init_state = inference_.init_prepare(model_inputs[0], key=init_key)
+            init_state = inference_.init_prepare(key=init_key)
             states = filter(
                 inference_,
                 model_inputs[1:],
@@ -297,7 +297,7 @@ def construct_chol_innovation_covariance(Y, chol_R, model_inputs):
             n_particles=n_particles,
             perturbed_obs=False,
         )
-        init_state = inference.init_prepare(model_inputs[0], key=init_key)
+        init_state = inference.init_prepare(key=init_key)
         states = filter(
             inference,
             model_inputs[1:],
@@ -336,7 +336,7 @@ def test_filter_noop(seed, x_dim, y_dim):
 
     inference, _ = load_enkf_inference(*lgssm, noop=True)
 
-    init_state = inference.init_prepare(jnp.array(0), key=random.key(seed + 1))
+    init_state = inference.init_prepare(key=random.key(seed + 1))
     prep_state = inference.filter_prepare(jnp.array(1), key=random.key(seed + 2))
     filtered_state = inference.filter_combine(init_state, prep_state)
 
@@ -364,7 +364,7 @@ def test_filter_noop(seed, x_dim, y_dim):
 def test_build_filter_requires_at_least_two_particles():
     """EnKF should fail fast when configured with fewer than two particles."""
 
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         return jnp.zeros(1) + jnp.eye(1) @ random.normal(key, (1,))
 
     with pytest.raises(ValueError, match="at least 2"):
diff --git a/tests/cuthbert/ensemble_kalman/test_enrts.py b/tests/cuthbert/ensemble_kalman/test_enrts.py
index 5e27bdf4..9764a65d 100644
--- a/tests/cuthbert/ensemble_kalman/test_enrts.py
+++ b/tests/cuthbert/ensemble_kalman/test_enrts.py
@@ -24,7 +24,7 @@ def load_enrts_inference(m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys):
     n_particles = 100_000
     x_dim = m0.shape[0]
 
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         return m0 + chol_P0 @ random.normal(key, m0.shape)
 
     def get_dynamics(model_inputs):
@@ -66,7 +66,7 @@ def test_smoother(self, seed, x_dim, y_dim, num_time_steps):
         )
 
         init_key, filter_key = random.split(random.key(seed + 1))
-        init_state = filter_obj.init_prepare(model_inputs[0], key=init_key)
+        init_state = filter_obj.init_prepare(key=init_key)
         filtered_states = filter(
             filter_obj,
             model_inputs[1:],
diff --git a/tests/cuthbert/factorial/test_discrete.py b/tests/cuthbert/factorial/test_discrete.py
index c5f850d4..4617a8a2 100644
--- a/tests/cuthbert/factorial/test_discrete.py
+++ b/tests/cuthbert/factorial/test_discrete.py
@@ -65,9 +65,6 @@ def generate_factorial_hmm(
 def build_factorial_discrete_filter(model_params):
     init_dist, trans_matrices, local_obs_lls, factorial_indices = model_params
 
-    def get_factorial_init_dist(model_inputs):
-        return init_dist
-
     def get_local_trans_matrix(model_inputs):
         inds = factorial_indices[model_inputs - 1]
         local_trans = trans_matrices[model_inputs - 1, inds]
@@ -76,9 +73,7 @@ def get_local_trans_matrix(model_inputs):
     def get_local_obs_lls(model_inputs):
         return local_obs_lls[model_inputs - 1]
 
-    filter_obj = build_filter(
-        get_factorial_init_dist, get_local_trans_matrix, get_local_obs_lls
-    )
+    filter_obj = build_filter(init_dist, get_local_trans_matrix, get_local_obs_lls)
     factorializer = factorial.discrete.build_factorializer(
         lambda model_inputs: factorial_indices[model_inputs - 1]
     )
@@ -160,7 +155,7 @@ def test_factorial_discrete_filter(
     )
 
     # output_factorial=False
-    init_state = filter_obj.init_prepare(model_inputs[0])
+    init_state = filter_obj.init_prepare()
     local_filter_states, final_state = factorial.filter(
         filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=False
     )
diff --git a/tests/cuthbert/factorial/test_filtering.py b/tests/cuthbert/factorial/test_filtering.py
index 8c43be2b..8343c768 100644
--- a/tests/cuthbert/factorial/test_filtering.py
+++ b/tests/cuthbert/factorial/test_filtering.py
@@ -31,9 +31,6 @@ def predict(mean, chol_cov, transition, offset, transition_chol_cov):
     def get_factorial_indices(x):
         return x
 
-    def get_init_params(_):
-        return means, chol_covs
-
     def get_dynamics_params(factor_index):
         return (
             transitions[factor_index],
@@ -46,10 +43,10 @@ def get_observation_params(_):
         return jnp.zeros((1, 1)), jnp.zeros(1), jnp.eye(1), jnp.array([jnp.nan])
 
     filter_obj = kalman.build_filter(
-        get_init_params, get_dynamics_params, get_observation_params
+        means, chol_covs, get_dynamics_params, get_observation_params
     )
     factorializer = factorial.gaussian.build_factorializer(get_factorial_indices)
-    factorial_state = filter_obj.init_prepare(None)
+    factorial_state = filter_obj.init_prepare()
 
     synchronized = factorial.synchronize(
         filter_obj, factorializer, jnp.arange(len(means)), factorial_state
@@ -66,14 +63,13 @@ def get_observation_params(_):
 def test_synchronize_discrete():
     init_dist = jnp.array([[0.6, 0.4], [0.25, 0.75]])
     transitions = jnp.array([[[0.7, 0.3], [0.2, 0.8]], [[0.4, 0.6], [0.9, 0.1]]])
-    initial_model_inputs = jnp.array(0)
     model_inputs = jnp.arange(len(init_dist))
     true_dists = vmap(lambda dist, transition: dist @ transition)(
         init_dist, transitions
     )
 
     filter_obj = build_discrete_filter(
-        lambda _: init_dist,
+        init_dist,
         lambda factor_index: transitions[factor_index],
         lambda _: jnp.zeros(2),
     )
@@ -87,7 +83,7 @@ def get_factorial_indices(x):
         filter_obj,
         factorializer,
         model_inputs,
-        filter_obj.init_prepare(initial_model_inputs),
+        filter_obj.init_prepare(),
     )
 
     chex.assert_trees_all_close(
@@ -97,11 +93,11 @@ def get_factorial_indices(x):
 
 
 def test_synchronize_particle_filter():
-    initial_model_inputs = jnp.array([1.0, 10.0])
+    initial_values = jnp.array([1.0, 10.0])
     model_inputs = jnp.array([0.5, -1.0])
 
     filter_obj = build_particle_filter(
-        init_sample=lambda key, initial_value: initial_value,
+        init_sample=lambda key: initial_values,
         propagate_sample=lambda key, particle, increment: particle + increment,
         log_potential=lambda previous, particle, increment: jnp.array(0.0),
         n_filter_particles=3,
@@ -115,8 +111,7 @@ def get_factorial_indices(x):
         get_factorial_indices, no_resampling.resampling
     )
     initial_state = factorializer.factorialize_init_state(
-        filter_obj.init_prepare(initial_model_inputs, key=random.key(0)),
-        initial_model_inputs,
+        filter_obj.init_prepare(key=random.key(0)),
     )
     true_particles = vmap(lambda particles, increment: particles + increment)(
         initial_state.particles, model_inputs
diff --git a/tests/cuthbert/factorial/test_kalman.py b/tests/cuthbert/factorial/test_kalman.py
index 0a8fbbf5..656acca6 100644
--- a/tests/cuthbert/factorial/test_kalman.py
+++ b/tests/cuthbert/factorial/test_kalman.py
@@ -47,9 +47,6 @@ def build_pairwise_factorial_filter(
 
     m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, factorial_indices = model_params
 
-    def get_init_params(model_inputs: int) -> tuple[Array, Array]:
-        return m0, chol_P0
-
     def get_dynamics_params(model_inputs: int) -> tuple[Array, Array, Array]:
         return Fs[model_inputs - 1], cs[model_inputs - 1], chol_Qs[model_inputs - 1]
 
@@ -62,7 +59,7 @@ def get_observation_params(model_inputs: int) -> tuple[Array, Array, Array, Arra
         )
 
     filter = kalman.build_filter(
-        get_init_params, get_dynamics_params, get_observation_params
+        m0, chol_P0, get_dynamics_params, get_observation_params
     )
 
     factorializer = factorial.gaussian.build_factorializer(
@@ -210,7 +207,7 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste
     fac_covs_t_all = jnp.stack(fac_covs_t_all)
 
     # Check output_factorial = False
-    init_state = filter_obj.init_prepare(model_inputs[0])
+    init_state = filter_obj.init_prepare()
     local_filter_states, final_state = factorial.filter(
         filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=False
     )
@@ -298,7 +295,7 @@ def test_smoother(
         smoother_factorial_index=smoother_factorial_index,
     )
 
-    init_state = filter_obj.init_prepare(filter_model_inputs[0])
+    init_state = filter_obj.init_prepare()
     local_filter_states, _ = factorial.filter(
         filter_obj,
         factorializer,
@@ -322,7 +319,8 @@ def test_smoother(
             factorializer.extract(init_state, smoother_factorial_index),
         )
         chex.assert_trees_all_close(
-            local_filter_states_single_factor, expected_initial_state
+            local_filter_states_single_factor._replace(model_inputs=None),
+            expected_initial_state,
         )
         return
 
diff --git a/tests/cuthbert/factorial/test_smc.py b/tests/cuthbert/factorial/test_smc.py
index 2c434644..1d5067b5 100644
--- a/tests/cuthbert/factorial/test_smc.py
+++ b/tests/cuthbert/factorial/test_smc.py
@@ -33,7 +33,7 @@ def build_factorial_smc_filter(
     m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, factorial_indices = model_params
     num_factors, x_dim = m0.shape
 
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         # Generates a single particle for all factors
         eps = random.normal(key, (num_factors, x_dim))
         return m0 + jax.vmap(lambda chol_P, e: chol_P @ e)(chol_P0, eps)
@@ -129,10 +129,8 @@ def test_factorial_smc_filter(
     kalman_filter, kalman_factorializer, kalman_model_inputs = (
         build_pairwise_factorial_filter(model_params)
     )
-    init_state = kalman_filter.init_prepare(kalman_model_inputs[0])
-    init_state = kalman_factorializer.factorialize_init_state(
-        init_state, kalman_model_inputs[0]
-    )
+    init_state = kalman_filter.init_prepare()
+    init_state = kalman_factorializer.factorialize_init_state(init_state)
     kalman_states = factorial.filter(
         kalman_filter,
         kalman_factorializer,
@@ -152,14 +150,11 @@ def test_factorial_smc_filter(
     smc_filter, smc_factorializer, smc_model_inputs = build_factorial_smc_filter(
         model_params, n_particles=num_particles
     )
-    smc_init_model_inputs = tree.map(lambda x: x[0], smc_model_inputs)
     smc_filter_model_inputs = tree.map(lambda x: x[1:], smc_model_inputs)
 
     init_key, filter_key = random.split(random.key(seed + 123))
-    init_smc_state = smc_filter.init_prepare(smc_init_model_inputs, key=init_key)
-    init_smc_state = smc_factorializer.factorialize_init_state(
-        init_smc_state, smc_init_model_inputs
-    )
+    init_smc_state = smc_filter.init_prepare(key=init_key)
+    init_smc_state = smc_factorializer.factorialize_init_state(init_smc_state)
     smc_states = factorial.filter(
         smc_filter,
         smc_factorializer,
diff --git a/tests/cuthbert/factorial/test_taylor.py b/tests/cuthbert/factorial/test_taylor.py
index f160ee05..1ce7401c 100644
--- a/tests/cuthbert/factorial/test_taylor.py
+++ b/tests/cuthbert/factorial/test_taylor.py
@@ -26,11 +26,10 @@ def test_factorial_taylor_filter_jit():
     factorial_indices = jnp.array([[0, 1], [1, 2], [0, 2]])
 
     # Users have to specify a initial log density that acts on the full factorial state.
-    def get_init_log_density(model_inputs):
-        def _init_log_density(x):
-            return jnp.sum(jax.vmap(multivariate_normal.logpdf)(x, m0s, chol_P0s))
+    def init_log_density(x):
+        return jnp.sum(jax.vmap(multivariate_normal.logpdf)(x, m0s, chol_P0s))
 
-        return _init_log_density, jnp.zeros((num_factors, x_dim))
+    init_linearization_point = jnp.zeros((num_factors, x_dim))
 
     def get_dynamics_log_density(state, model_inputs):
         F = Fs[model_inputs - 1]
@@ -53,18 +52,18 @@ def observation_log_density(x, y):
         return observation_log_density, state.mean, ys[model_inputs - 1]
 
     filter_obj = taylor.build_filter(
-        get_init_log_density,
+        init_log_density,
+        init_linearization_point,
         get_dynamics_log_density,
         get_observation_log_density,
     )
     factorializer = factorial.gaussian.build_factorializer(
         lambda model_inputs: factorial_indices[model_inputs - 1]
     )
-    init_model_inputs = jnp.array(0)
     filter_model_inputs = jnp.arange(1, num_time_steps + 1)
 
-    init_state = filter_obj.init_prepare(init_model_inputs)
-    init_state = factorializer.factorialize_init_state(init_state, init_model_inputs)
+    init_state = filter_obj.init_prepare()
+    init_state = factorializer.factorialize_init_state(init_state)
     local_filter_states, _ = jax.jit(
         factorial.filter,
         static_argnames=("filter_obj", "factorializer", "output_factorial"),
diff --git a/tests/cuthbert/factorial/test_utils.py b/tests/cuthbert/factorial/test_utils.py
index 3fec75a7..78633e2d 100644
--- a/tests/cuthbert/factorial/test_utils.py
+++ b/tests/cuthbert/factorial/test_utils.py
@@ -133,3 +133,25 @@ def test_serial_to_single_factor_prepends_init_factorial_tree():
     )
 
     chex.assert_trees_all_close(factor_1, all_factors[factorial_index])
+
+
+def test_serial_to_factorial_pads_missing_initial_subtree_without_occurrences():
+    def extract_state(state, inds):
+        return {"x": state["x"][inds], "inputs": state["inputs"]}
+
+    serial_tree = {
+        "x": jnp.array([[[2.0]]]),
+        "inputs": {"y": jnp.array([[3.0]]), "mask": jnp.array([True])},
+    }
+    initial_tree = {"x": jnp.array([[1.0], [4.0]]), "inputs": None}
+    result = serial_to_factorial(
+        extract_state,
+        serial_tree,
+        jnp.array([[0]]),
+        select_factorial_inds=1,
+        init_factorial_tree=initial_tree,
+    )
+    chex.assert_trees_all_equal(result["x"], jnp.array([[4.0]]))
+    chex.assert_shape(result["inputs"]["y"], (1, 1))
+    assert result["inputs"]["mask"].dtype == jnp.bool_
+    assert initial_tree["inputs"] is None
diff --git a/tests/cuthbert/gaussian/test_kalman.py b/tests/cuthbert/gaussian/test_kalman.py
index 21501598..8918b9b1 100644
--- a/tests/cuthbert/gaussian/test_kalman.py
+++ b/tests/cuthbert/gaussian/test_kalman.py
@@ -69,9 +69,6 @@ def load_kalman_inference(
 ) -> tuple[Filter, Smoother, Array]:
     """Builds Kalman filter and smoother objects and model_inputs for a linear-Gaussian SSM."""
 
-    def get_init_params(model_inputs: int) -> tuple[Array, Array]:
-        return m0, chol_P0
-
     def get_dynamics_params(model_inputs: int) -> tuple[Array, Array, Array]:
         return Fs[model_inputs - 1], cs[model_inputs - 1], chol_Qs[model_inputs - 1]
 
@@ -84,7 +81,7 @@ def get_observation_params(model_inputs: int) -> tuple[Array, Array, Array, Arra
         )
 
     filter = kalman.build_filter(
-        get_init_params, get_dynamics_params, get_observation_params
+        m0, chol_P0, get_dynamics_params, get_observation_params
     )
     smoother = kalman.build_smoother(
         get_dynamics_params, store_gain=True, store_chol_cov_given_next=True
@@ -119,7 +116,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
         m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys
     )
 
-    init_state = kalman_filter.init_prepare(model_inputs[0])
+    init_state = kalman_filter.init_prepare()
 
     # Run sequential sqrt filter
     seq_states = filter(kalman_filter, model_inputs[1:], init_state, parallel=False)
@@ -168,7 +165,7 @@ def get_loglikelihood_from_params(
         kalman_filter, _, model_inputs = load_kalman_inference(
             m0_, chol_P0_, Fs_, cs_, chol_Qs_, Hs_, ds_, chol_Rs_, ys
         )
-        init_state = kalman_filter.init_prepare(model_inputs[0])
+        init_state = kalman_filter.init_prepare()
         states = filter(kalman_filter, model_inputs[1:], init_state)
         return states.log_normalizing_constant[-1]
 
@@ -189,7 +186,8 @@ def test_filter_noop(seed, x_dim, y_dim):
     m0, chol_P0 = generate_lgssm(seed, x_dim, y_dim, 0)[:2]
 
     filter_obj = kalman.build_filter(
-        get_init_params=lambda model_inputs: (m0, chol_P0),
+        m0=m0,
+        chol_P0=chol_P0,
         get_dynamics_params=lambda model_inputs: (  # p(x_t | x_{t-1}) = N(x_t | x_{t-1}, 0)
             jnp.eye(x_dim),
             jnp.zeros(x_dim),
@@ -207,7 +205,7 @@ def test_filter_noop(seed, x_dim, y_dim):
         ),
     )
 
-    state = filter_obj.init_prepare(None)
+    state = filter_obj.init_prepare()
     prep_state = filter_obj.filter_prepare(None)
     filtered_state = filter_obj.filter_combine(state, prep_state)
 
@@ -234,7 +232,7 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    init_state = kalman_filter.init_prepare(model_inputs[0])
+    init_state = kalman_filter.init_prepare()
     filt_states = filter(kalman_filter, model_inputs[1:], init_state)
     filt_means, filt_chol_covs = filt_states.mean, filt_states.chol_cov
     filt_covs = filt_chol_covs @ filt_chol_covs.transpose(0, 2, 1)
diff --git a/tests/cuthbert/gaussian/test_moments.py b/tests/cuthbert/gaussian/test_moments.py
index 22a2150c..d89e543b 100644
--- a/tests/cuthbert/gaussian/test_moments.py
+++ b/tests/cuthbert/gaussian/test_moments.py
@@ -35,9 +35,6 @@ def load_moments_inference(
 ) -> tuple[Filter, Smoother, Array]:
     """Builds linearized moments Kalman filter and smoother objects and model_inputs for a linear-Gaussian SSM."""
 
-    def get_init_params(model_inputs: int) -> tuple[Array, Array]:
-        return m0, chol_P0
-
     def dynamics_moments(state, model_inputs):
         def dynamics_mean_and_chol_cov_func(x):
             return Fs[model_inputs - 1] @ x + cs[model_inputs - 1], chol_Qs[
@@ -60,7 +57,8 @@ def observation_mean_and_chol_cov_and_y_func(x):
         )
 
     filter = moments.build_filter(
-        get_init_params,
+        m0,
+        chol_P0,
         dynamics_moments,
         observation_moments,
         associative=associative_filter,
@@ -93,7 +91,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
         m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, associative_filter=False
     )
 
-    init_state = moments_filter.init_prepare(model_inputs[0])
+    init_state = moments_filter.init_prepare()
 
     # Run sequential sqrt filter
     seq_states = filter(moments_filter, model_inputs[1:], init_state, parallel=False)
@@ -107,7 +105,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
         m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, associative_filter=True
     )
 
-    associative_init_state = associative_moments_filter.init_prepare(model_inputs[0])
+    associative_init_state = associative_moments_filter.init_prepare()
 
     # Run associative filter with parallel=False
     seq_ass_states = filter(
@@ -181,13 +179,14 @@ def observation_mean_and_chol_cov_and_y_func(x):
         )
 
     filter_obj = moments.build_filter(
-        get_init_params=lambda model_inputs: (m0, chol_P0),
+        m0=m0,
+        chol_P0=chol_P0,
         get_dynamics_params=noop_dynamics_moments,
         get_observation_params=noop_observation_moments,
         associative=associative,
     )
 
-    state = filter_obj.init_prepare(None)
+    state = filter_obj.init_prepare()
     prep_state = filter_obj.filter_prepare(None)
     filtered_state = filter_obj.filter_combine(state, prep_state)
 
@@ -214,7 +213,7 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    init_state = extended_filter.init_prepare(model_inputs[0])
+    init_state = extended_filter.init_prepare()
     filt_states = filter(extended_filter, model_inputs[1:], init_state)
     filt_means, filt_chol_covs = filt_states.mean, filt_states.chol_cov
     filt_covs = filt_chol_covs @ filt_chol_covs.transpose(0, 2, 1)
diff --git a/tests/cuthbert/gaussian/test_taylor.py b/tests/cuthbert/gaussian/test_taylor.py
index 7ca06960..7ce23619 100644
--- a/tests/cuthbert/gaussian/test_taylor.py
+++ b/tests/cuthbert/gaussian/test_taylor.py
@@ -10,7 +10,6 @@
 from cuthbert.gaussian import taylor
 from cuthbert.gaussian.taylor.types import (
     GetDynamicsLogDensity,
-    GetInitLogDensity,
     LogConditionalDensity,
     LogDensity,
     LogPotential,
@@ -38,16 +37,15 @@ def _load_taylor_init_and_dynamics(
     Fs: Array,
     cs: Array,
     chol_Qs: Array,
-) -> tuple[GetInitLogDensity, GetDynamicsLogDensity]:
+) -> tuple[tuple[LogDensity, Array], GetDynamicsLogDensity]:
     """Builds linearized log density Kalman filter and smoother objects and model_inputs
     for a linear-Gaussian SSM.
     """
 
-    def get_init_log_density(model_inputs: int) -> tuple[LogDensity, Array]:
-        def init_log_density(x):
-            return multivariate_normal.logpdf(x, m0, chol_P0)
+    def init_log_density(x):
+        return multivariate_normal.logpdf(x, m0, chol_P0)
 
-        return init_log_density, jnp.zeros_like(m0)
+    init_linearization_point = jnp.zeros_like(m0)
 
     def get_dynamics_log_density(
         state: LinearizedKalmanFilterState, model_inputs: int
@@ -65,7 +63,7 @@ def dynamics_log_density(x_prev, x):
             jnp.zeros_like(m0),
         )
 
-    return get_init_log_density, get_dynamics_log_density
+    return (init_log_density, init_linearization_point), get_dynamics_log_density
 
 
 def load_taylor_inference(
@@ -84,7 +82,7 @@ def load_taylor_inference(
     """Builds linearized log density Kalman filter and smoother objects and model_inputs
     for a linear-Gaussian SSM.
     """
-    get_init_log_density, get_dynamics_log_density = _load_taylor_init_and_dynamics(
+    init_params, get_dynamics_log_density = _load_taylor_init_and_dynamics(
         m0, chol_P0, Fs, cs, chol_Qs
     )
 
@@ -105,7 +103,7 @@ def observation_log_density(x, y):
         )
 
     filter = taylor.build_filter(
-        get_init_log_density,
+        *init_params,
         get_dynamics_log_density,
         get_observation_log_density,
         associative=associative_filter,
@@ -153,7 +151,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
         ignore_nan_dims=True,
     )
 
-    init_state = taylor_filter.init_prepare(model_inputs[0])
+    init_state = taylor_filter.init_prepare()
 
     # Run sequential sqrt filter
     seq_states = filter(taylor_filter, model_inputs[1:], init_state, parallel=False)
@@ -177,7 +175,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
         ignore_nan_dims=True,
     )
 
-    associative_init_state = associative_taylor_filter.init_prepare(model_inputs[0])
+    associative_init_state = associative_taylor_filter.init_prepare()
 
     # Run associative filter with parallel=False§
     seq_ass_states = filter(
@@ -233,11 +231,10 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
 def test_filter_noop(seed, x_dim, y_dim, associative):
     m0, chol_P0 = generate_lgssm(seed, x_dim, y_dim, 0)[:2]
 
-    def get_init_log_density(model_inputs: int) -> tuple[LogDensity, Array]:
-        def init_log_density(x):
-            return multivariate_normal.logpdf(x, m0, chol_P0)
+    def init_log_density(x):
+        return multivariate_normal.logpdf(x, m0, chol_P0)
 
-        return init_log_density, jnp.zeros_like(m0)
+    init_linearization_point = jnp.zeros_like(m0)
 
     def get_noop_dynamics_log_density(
         state: LinearizedKalmanFilterState, model_inputs: int
@@ -263,13 +260,14 @@ def get_noop_observation_log_density(
         )
 
     filter_obj = taylor.build_filter(
-        get_init_log_density=get_init_log_density,
+        init_log_density,
+        init_linearization_point,
         get_dynamics_log_density=get_noop_dynamics_log_density,
         get_observation_func=get_noop_observation_log_density,
         associative=associative,
     )
 
-    state = filter_obj.init_prepare(None)
+    state = filter_obj.init_prepare()
     prep_state = filter_obj.filter_prepare(None)
     filtered_state = filter_obj.filter_combine(state, prep_state)
 
@@ -293,7 +291,8 @@ def get_noop_observation_log_potential(
         )
 
     filter_log_potential_obj = taylor.build_filter(
-        get_init_log_density=get_init_log_density,
+        init_log_density,
+        init_linearization_point,
         get_dynamics_log_density=get_noop_dynamics_log_density,
         get_observation_func=get_noop_observation_log_potential,
         associative=associative,
@@ -328,7 +327,7 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    init_state = log_density_filter.init_prepare(model_inputs[0])
+    init_state = log_density_filter.init_prepare()
     filt_states = filter(log_density_filter, model_inputs[1:], init_state)
     filt_means, filt_chol_covs = filt_states.mean, filt_states.chol_cov
     filt_covs = filt_chol_covs @ filt_chol_covs.transpose(0, 2, 1)
@@ -391,7 +390,7 @@ def load_taylor_inference_potential(
 
     Uses Gaussian log potential for observations instead of conditional log density.
     """
-    get_init_log_density, get_dynamics_log_density = _load_taylor_init_and_dynamics(
+    init_params, get_dynamics_log_density = _load_taylor_init_and_dynamics(
         m0, chol_P0, Fs, cs, chol_Qs
     )
 
@@ -409,7 +408,7 @@ def observation_log_potential(x):
         )
 
     filter = taylor.build_filter(
-        get_init_log_density,
+        *init_params,
         get_dynamics_log_density,
         get_observation_log_potential,
         associative=associative_filter,
@@ -436,7 +435,7 @@ def test_offline_filter_potential(seed, x_dim, num_time_steps):
         m0, chol_P0, Fs, cs, chol_Qs, ms, chol_Rs, associative_filter=False
     )
 
-    init_state = taylor_filter.init_prepare(model_inputs[0])
+    init_state = taylor_filter.init_prepare()
 
     # Run sequential sqrt filter
     seq_states = filter(taylor_filter, model_inputs[1:], init_state, parallel=False)
@@ -450,7 +449,7 @@ def test_offline_filter_potential(seed, x_dim, num_time_steps):
         m0, chol_P0, Fs, cs, chol_Qs, ms, chol_Rs, associative_filter=True
     )
 
-    associative_init_state = associative_taylor_filter.init_prepare(model_inputs[0])
+    associative_init_state = associative_taylor_filter.init_prepare()
 
     # Run associative filter with parallel=False
     seq_ass_states = filter(
@@ -515,18 +514,19 @@ def test_factorial_init(seed, x_dim, num_factors):
     chol_P0 = jnp.linalg.cholesky(P0)
     lin_p = random.normal(lin_p_key, shape=(num_factors, x_dim))
 
-    def get_init_log_density(model_inputs):
-        def init_log_density(x):
-            factor_log_densities = jax.vmap(
-                lambda factor_x: multivariate_normal.logpdf(factor_x, m0, chol_P0)
-            )(x)
-            return jnp.sum(factor_log_densities)
+    def init_log_density(x):
+        factor_log_densities = jax.vmap(
+            lambda factor_x: multivariate_normal.logpdf(factor_x, m0, chol_P0)
+        )(x)
+        return jnp.sum(factor_log_densities)
 
-        return init_log_density, lin_p
+    init_linearization_point = lin_p
 
-    init_state = taylor.non_associative_filter.init_prepare(None, get_init_log_density)
+    init_state = taylor.non_associative_filter.init_prepare(
+        init_log_density, init_linearization_point
+    )
     init_state_assoc = taylor.associative_filter.init_prepare(
-        None, get_init_log_density
+        init_log_density, init_linearization_point
     )
 
     factorial_mean = jnp.tile(m0[None], (num_factors, 1))
diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py
index fbf9eabc..efdaadd3 100644
--- a/tests/cuthbert/smc/test_backward_sampler.py
+++ b/tests/cuthbert/smc/test_backward_sampler.py
@@ -28,7 +28,7 @@ def config():
 
 
 def load_inference(m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys):
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         return m0 + chol_P0 @ random.normal(key, m0.shape)
 
     def propagate_sample(key, state, model_inputs: int):
@@ -76,7 +76,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
             m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys
         )
         init_key, filter_key, smoother_key = random.split(random.key(seed + 1), 3)
-        init_state = filter_obj.init_prepare(model_inputs[0], key=init_key)
+        init_state = filter_obj.init_prepare(key=init_key)
         filtered_states = filter(
             filter_obj,
             model_inputs[1:],
@@ -124,7 +124,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
     def test_pytree_particles(self, method):
         """Test that the pf handles pytree states correctly."""
 
-        def init_sample(key, model_inputs):
+        def init_sample(key):
             keys = random.split(key, 2)
             position = random.normal(keys[0], (2,))
             velocity = random.normal(keys[1], (2,))
@@ -164,7 +164,7 @@ def log_potential(state_prev, state, model_inputs):
         init_key, filter_key, smoother_key = random.split(key, 3)
         num_time_steps = 5
         model_inputs = jnp.empty(num_time_steps + 1)
-        init_state = filter_obj.init_prepare(model_inputs[0], key=init_key)
+        init_state = filter_obj.init_prepare(key=init_key)
         filtered_states = filter(
             filter_obj,
             model_inputs[1:],
diff --git a/tests/cuthbert/smc/test_particle_filters.py b/tests/cuthbert/smc/test_particle_filters.py
index 3de3286a..a9db5bfb 100644
--- a/tests/cuthbert/smc/test_particle_filters.py
+++ b/tests/cuthbert/smc/test_particle_filters.py
@@ -37,7 +37,7 @@ def load_inference(
     else:
         raise ValueError(f"Unknown method: {method}")
 
-    def init_sample(key, model_inputs):
+    def init_sample(key):
         return m0 + chol_P0 @ random.normal(key, m0.shape)
 
     if noop:
@@ -108,7 +108,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
             m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, method
         )
         init_key, filter_key = random.split(random.key(seed + 1))
-        init_state = inference.init_prepare(model_inputs[0], key=init_key)
+        init_state = inference.init_prepare(key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
             inference,
             model_inputs[1:],
@@ -149,7 +149,7 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
     def test_pytree_particles(self, method):
         """Test that the pf handles pytree states correctly."""
 
-        def init_sample(key, model_inputs):
+        def init_sample(key):
             keys = random.split(key, 2)
             position = random.normal(keys[0], (2,))
             velocity = random.normal(keys[1], (2,))
@@ -204,7 +204,7 @@ def log_potential(state_prev, state, model_inputs):
 
         # Run the particle filter
         init_key, filter_key = random.split(key)
-        init_state = inference.init_prepare(model_inputs[0], key=init_key)
+        init_state = inference.init_prepare(key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
             inference,
             model_inputs[1:],
@@ -229,7 +229,7 @@ def test_filter_noop(seed, x_dim, y_dim, method):
 
     inference, _ = load_inference(*lgssm, method=method, noop=True)
 
-    init_state = inference.init_prepare(None, key=random.key(seed + 1))
+    init_state = inference.init_prepare(key=random.key(seed + 1))
     prep_state = inference.filter_prepare(None, key=random.key(seed + 2))
     filtered_state = inference.filter_combine(init_state, prep_state)