diff --git a/cuthbert/README.md b/cuthbert/README.md
index 8602bf80..418139f7 100644
--- a/cuthbert/README.md
+++ b/cuthbert/README.md
@@ -10,7 +10,8 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-model_inputs = ...
+init_model_inputs = ...
+filter_model_inputs = ...
 
 # Load inference method
 kalman_filter = cuthbert.gaussian.kalman.build_filter(
@@ -20,10 +21,10 @@ kalman_filter = cuthbert.gaussian.kalman.build_filter(
 )   # build_filter function takes all inference-specific arguments, swap this out for different inference methods.
 
 # Online inference
-state = kalman_filter.init_prepare(tree.map(lambda x: x[0], model_inputs))
+state = kalman_filter.init_prepare(init_model_inputs)
 
-for t in range(1, T):
-    model_inputs_t = tree.map(lambda x: x[t], model_inputs)
+for t in range(T):
+    model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
     prepare_state = kalman_filter.filter_prepare(model_inputs_t)
     state = kalman_filter.filter_combine(state, prepare_state)
 ```
@@ -33,7 +34,8 @@ Or for offline inference:
 ```python
 kalman_smoother = cuthbert.gaussian.kalman.build_smoother(get_dynamics_params)
 
-filter_states = cuthbert.filter(kalman_filter, model_inputs)
-smoother_states = cuthbert.smoother(kalman_smoother, filter_states, model_inputs)
+init_state = kalman_filter.init_prepare(init_model_inputs)
+filter_states = cuthbert.filter(kalman_filter, filter_model_inputs, init_state)
+smoother_states = cuthbert.smoother(kalman_smoother, filter_states, filter_model_inputs)
 ```
 <!-- --8<-- [end:unified_interface] -->
diff --git a/cuthbert/factorial/README.md b/cuthbert/factorial/README.md
index df3246e1..898df3b4 100644
--- a/cuthbert/factorial/README.md
+++ b/cuthbert/factorial/README.md
@@ -47,7 +47,8 @@ from jax import tree
 import cuthbert
 
 # Define model_inputs
-model_inputs = ...
+init_model_inputs = ...
+filter_model_inputs = ...
 
 # Define function to extract the factorial indices from model inputs
 # Here we assume model_inputs is a NamedTuple with a field `factorial_inds`
@@ -64,12 +65,11 @@ kalman_filter = cuthbert.gaussian.kalman.build_filter(
 )
 
 # Online inference
-init_model_inputs = tree.map(lambda x: x[0], model_inputs)
 factorial_state = kalman_filter.init_prepare(init_model_inputs)
 factorial_state = factorializer.factorialize_init_state(factorial_state, init_model_inputs)
 
-for t in range(1, T):
-    model_inputs_t = tree.map(lambda x: x[t], model_inputs)
+for t in range(T):
+    model_inputs_t = tree.map(lambda x: x[t], filter_model_inputs)
     local_state = factorializer.extract_and_join(prev_factorial_state, model_inputs_t)
     prepare_state = kalman_filter.filter_prepare(model_inputs_t)
     local_joint_filtered_state = kalman_filter.filter_combine(local_state, prepare_state)
@@ -82,9 +82,12 @@ 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, local_filter_states, final_factorial_state = (
+init_factorial_state = kalman_filter.init_prepare(init_model_inputs)
+init_factorial_state = factorializer.factorialize_init_state(init_factorial_state, init_model_inputs)
+
+local_filter_states, final_factorial_state = (
     cuthbert.factorial.filter(
-        kalman_filter, factorializer, model_inputs, output_factorial=False
+        kalman_filter, factorializer, filter_model_inputs, init_factorial_state, output_factorial=False
     )
 )
 ```
diff --git a/cuthbert/factorial/filtering.py b/cuthbert/factorial/filtering.py
index cc02ae38..2ff4c624 100644
--- a/cuthbert/factorial/filtering.py
+++ b/cuthbert/factorial/filtering.py
@@ -13,16 +13,21 @@ def filter(
     filter_obj: Filter,
     factorializer: Factorializer,
     model_inputs: ArrayTreeLike,
+    init_state: ArrayTree,
+    *,
     output_factorial: bool = False,
     key: KeyArray | None = None,
 ) -> (
-    ArrayTree | tuple[ArrayTree, ArrayTree, ArrayTree]
+    ArrayTree | tuple[ArrayTree, ArrayTree]
 ):  # TODO: Can overload this function so the type checker knows that the output is an ArrayTree if output_factorial is True and a tuple[ArrayTree, ArrayTree, ArrayTree] if output_factorial is False
     """Applies offline factorial filtering for given model inputs.
 
-    `model_inputs` should have leading temporal dimension of length T + 1,
+    `model_inputs` should have leading temporal dimension of length T,
     where T is the number of time steps excluding the initial state.
 
+    The output will have leading temporal dimension of length T + 1, and
+    include the initial state as the first element.
+
     Parallel associative filtering is not supported for factorial filtering.
 
     Note that if output_factorial is True, this function will output a factorial state
@@ -33,14 +38,18 @@ def filter(
     Args:
         filter_obj: The filter inference object.
         factorializer: The factorializer object for the inference method.
-        model_inputs: The model inputs (with leading temporal dimension of length T + 1).
+        model_inputs: The model inputs for filtering
+            (with leading temporal dimension of length T).
+        init_state: The initial factorial state
+            (with factorial dimension but no temporal dimension).
+            Generated by `filter_obj.init_prepare` and then factorialized by
+            `factorializer.factorialize_init_state`.
         output_factorial: If True, return a single state with first temporal dimension
             of length T + 1 and second factorial dimension of length F.
-            If False, return a tuple of states. The first being the initial state
-            with first dimension of length F and no temporal dimension.
-            The second being the local states for each time step, i.e. first
+            If False, return a tuple of states.
+            The first being the local states for each time step, i.e. first
             dimension of length T and no factorial dimension.
-            The third being the final factorial state with first dimension of
+            The second being the final factorial state with first dimension of
             length F and no temporal dimension.
         key: The key for the random number generator.
 
@@ -48,27 +57,16 @@ def filter(
         If output_factorial is True, returns a single state with first temporal dimension
             of length T + 1 and second factorial dimension of length F.
         If output_factorial is False, returns a tuple of
-            (initial factorial state, local states for each time step,
-            final factorial state).
+            (local states for each time step, final factorial state).
     """
-    T = tree.leaves(model_inputs)[0].shape[0] - 1
+    T = tree.leaves(model_inputs)[0].shape[0]
 
     if key is None:
         # This will throw error if used as a key, which is desired behavior
         # (albeit not a useful error, we could improve this)
-        prepare_keys = jnp.empty(T + 1)
+        prepare_keys = jnp.empty(T)
     else:
-        prepare_keys = random.split(key, T + 1)
-
-    init_model_input = tree.map(lambda x: x[0], model_inputs)
-    init_factorial_state = filter_obj.init_prepare(
-        init_model_input, key=prepare_keys[0]
-    )
-    init_factorial_state = factorializer.factorialize_init_state(
-        init_factorial_state, init_model_input
-    )
-
-    prep_model_inputs = tree.map(lambda x: x[1:], model_inputs)
+        prepare_keys = random.split(key, T)
 
     def body_local(prev_factorial_state, prep_inp_and_k):
         prep_inp, k = prep_inp_and_k
@@ -99,12 +97,12 @@ def body_factorial(prev_factorial_state, prep_inp_and_k):
 
         _, factorial_states = scan(
             body_factorial,
-            init_factorial_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            init_state,
+            (model_inputs, prepare_keys),
         )
         factorial_states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]),
-            init_factorial_state,
+            init_state,
             factorial_states,
         )
 
@@ -113,7 +111,7 @@ def body_factorial(prev_factorial_state, prep_inp_and_k):
     else:
         final_factorial_state, local_states = scan(
             body_local,
-            init_factorial_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            init_state,
+            (model_inputs, prepare_keys),
         )
-        return init_factorial_state, local_states, final_factorial_state
+        return local_states, final_factorial_state
diff --git a/cuthbert/filtering.py b/cuthbert/filtering.py
index ce5c6561..cd1cff47 100644
--- a/cuthbert/filtering.py
+++ b/cuthbert/filtering.py
@@ -13,17 +13,25 @@
 def filter(
     filter_obj: Filter,
     model_inputs: ArrayTreeLike,
+    init_state: ArrayTree,
+    *,
     parallel: bool = False,
     key: KeyArray | None = None,
 ) -> ArrayTree:
     """Applies offline filtering given a filter object and model inputs.
 
-    `model_inputs` should have leading temporal dimension of length T + 1,
+    `model_inputs` should have leading temporal dimension of length T,
     where T is the number of time steps excluding the initial state.
 
+    The output will have leading temporal dimension of length T + 1, and
+    include the initial state as the first element.
+
     Args:
         filter_obj: The filter inference object.
-        model_inputs: The model inputs (with leading temporal dimension of length T + 1).
+        model_inputs: The model inputs for filtering
+            (with leading temporal dimension of length T).
+        init_state: The initial state (with no temporal dimension).
+            Generated by `filter_obj.init_prepare`.
         parallel: Whether to run the filter in parallel.
             Requires `filter.associative_filter` to be `True`.
         key: The key for the random number generator.
@@ -36,23 +44,18 @@ def filter(
             f"Parallel filtering attempted but filter.associative is False for {filter_obj}"
         )
 
-    T = tree.leaves(model_inputs)[0].shape[0] - 1
+    T = tree.leaves(model_inputs)[0].shape[0]
 
     if key is None:
         # This will throw error if used as a key, which is desired behavior
         # (albeit not a useful error, we could improve this)
-        prepare_keys = jnp.empty(T + 1)
+        prepare_keys = jnp.empty(T)
     else:
-        prepare_keys = random.split(key, T + 1)
-
-    init_model_input = tree.map(lambda x: x[0], model_inputs)
-    init_state = filter_obj.init_prepare(init_model_input, key=prepare_keys[0])
-
-    prep_model_inputs = tree.map(lambda x: x[1:], model_inputs)
+        prepare_keys = random.split(key, T)
 
     if parallel:
         other_prep_states = vmap(lambda inp, k: filter_obj.filter_prepare(inp, key=k))(
-            prep_model_inputs, prepare_keys[1:]
+            model_inputs, prepare_keys
         )
         prep_states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]), init_state, other_prep_states
@@ -72,7 +75,7 @@ def body(prev_state, prep_inp_and_k):
         _, states = scan(
             body,
             init_state,
-            (prep_model_inputs, prepare_keys[1:]),
+            (model_inputs, prepare_keys),
         )
         states = tree.map(
             lambda x, y: jnp.concatenate([x[None], y]), init_state, states
diff --git a/docs/assets/enkf_comparison.png b/docs/assets/enkf_comparison.png
index d9ab1533..5538d8b8 100644
Binary files a/docs/assets/enkf_comparison.png and b/docs/assets/enkf_comparison.png differ
diff --git a/docs/assets/enkf_comparison_loglik.png b/docs/assets/enkf_comparison_loglik.png
index f6a80a71..b1bb016f 100644
Binary files a/docs/assets/enkf_comparison_loglik.png and b/docs/assets/enkf_comparison_loglik.png differ
diff --git a/docs/assets/enkf_comparison_rmse.png b/docs/assets/enkf_comparison_rmse.png
index 131c68ae..dca6ca85 100644
Binary files a/docs/assets/enkf_comparison_rmse.png and b/docs/assets/enkf_comparison_rmse.png differ
diff --git a/docs/assets/enkf_comparison_timing.png b/docs/assets/enkf_comparison_timing.png
index 7a5ca678..188ddae3 100644
Binary files a/docs/assets/enkf_comparison_timing.png and b/docs/assets/enkf_comparison_timing.png differ
diff --git a/docs/assets/pf_grad_bias_simulated_data.png b/docs/assets/pf_grad_bias_simulated_data.png
index 10478530..43814510 100644
Binary files a/docs/assets/pf_grad_bias_simulated_data.png and b/docs/assets/pf_grad_bias_simulated_data.png differ
diff --git a/docs/examples/diff_pf_resampling.md b/docs/examples/diff_pf_resampling.md
index 2747223c..92be34b3 100644
--- a/docs/examples/diff_pf_resampling.md
+++ b/docs/examples/diff_pf_resampling.md
@@ -123,7 +123,8 @@ def kalman_mll(theta, ys, *, m0, chol_P0):
 
     kf = kalman.build_filter(get_init_params, get_dynamics_params, get_observation_params)
     model_inputs = jnp.arange(ys.shape[0] + 1)
-    states = run_filter(kf, model_inputs, parallel=False)
+    init_state = kf.init_prepare(model_inputs[0])
+    states = run_filter(kf, model_inputs[1:], init_state, parallel=False)
     return states.log_normalizing_constant[-1]
 
 
@@ -196,7 +197,11 @@ def pf_mll(
     )
 
     model_inputs = jnp.arange(ys.shape[0] + 1)
-    states = run_filter(filt, model_inputs, parallel=False, key=key)
+    init_key, filter_key = jax.random.split(key)
+    init_state = filt.init_prepare(model_inputs[0], key=init_key)
+    states = run_filter(
+        filt, model_inputs[1:], init_state, parallel=False, key=filter_key
+    )
     return states.log_normalizing_constant[-1]
 
 
@@ -536,7 +541,11 @@ Finally, we should be sure that the forward pass is not actually being modified
         )
 
         model_inputs = jnp.arange(ys.shape[0] + 1)
-        states = run_filter(filt, model_inputs, parallel=False, key=key)
+        init_key, filter_key = jax.random.split(key)
+        init_state = filt.init_prepare(model_inputs[0], key=init_key)
+        states = run_filter(
+            filt, model_inputs[1:], init_state, parallel=False, key=filter_key
+        )
         return states.particles[-1]
 
 
diff --git a/docs/examples/dynamax_integration.md b/docs/examples/dynamax_integration.md
index 0dfb47d4..1db24e18 100644
--- a/docs/examples/dynamax_integration.md
+++ b/docs/examples/dynamax_integration.md
@@ -171,7 +171,8 @@ Now we can run the `cuthbert` Kalman filter to obtain the filtering distribution
 
 ```{.python #dynamax-run-filter}
 # Run Kalman filtering
-filtered_states = cuthbert.filter(filter_obj, model_inputs)
+init_state = filter_obj.init_prepare(model_inputs[0])
+filtered_states = cuthbert.filter(filter_obj, model_inputs[1:], init_state)
 
 # Extract filtering results - remove initial time step
 filtered_means = filtered_states.mean[1:]
@@ -327,4 +328,3 @@ filters.
 <<dynamax-visualize>>
 ```
 -->
-
diff --git a/docs/examples/enkf_comparison.md b/docs/examples/enkf_comparison.md
index 8c8dcd4b..8c0f2853 100644
--- a/docs/examples/enkf_comparison.md
+++ b/docs/examples/enkf_comparison.md
@@ -164,14 +164,16 @@ ekf = taylor.build_filter(
 )
 
 jitted_filter = jit(run_filter, static_argnames=["filter_obj"])
+filter_model_inputs = model_inputs[1:]
 
 n_timing = 20
-ekf_states = jitted_filter(ekf, model_inputs)  # warm up
+ekf_init_state = ekf.init_prepare(model_inputs[0])
+ekf_states = jitted_filter(ekf, filter_model_inputs, ekf_init_state)  # warm up
 jax.block_until_ready(ekf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(ekf, model_inputs))
+    jax.block_until_ready(jitted_filter(ekf, filter_model_inputs, ekf_init_state))
     _times.append(time.perf_counter() - _t0)
 ekf_time = float(jnp.median(jnp.array(_times)))
 
@@ -194,13 +196,18 @@ enkf = ensemble_kalman_filter.build_filter(
     perturbed_obs=True,
 )
 
-key, enkf_key = random.split(key)
-enkf_states = jitted_filter(enkf, model_inputs, key=enkf_key)  # warm up
+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_states = jitted_filter(
+    enkf, filter_model_inputs, enkf_init_state, key=enkf_filter_key
+)  # warm up
 jax.block_until_ready(enkf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(enkf, model_inputs, key=enkf_key))
+    jax.block_until_ready(
+        jitted_filter(enkf, filter_model_inputs, enkf_init_state, key=enkf_filter_key)
+    )
     _times.append(time.perf_counter() - _t0)
 enkf_time = float(jnp.median(jnp.array(_times)))
 
@@ -225,13 +232,18 @@ pf = particle_filter.build_filter(
     resampling_fn=adaptive_systematic,
 )
 
-key, pf_key = random.split(key)
-pf_states = jitted_filter(pf, model_inputs, key=pf_key)  # warm up
+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_states = jitted_filter(
+    pf, filter_model_inputs, pf_init_state, key=pf_filter_key
+)  # warm up
 jax.block_until_ready(pf_states)
 _times = []
 for _ in range(n_timing):
     _t0 = time.perf_counter()
-    jax.block_until_ready(jitted_filter(pf, model_inputs, key=pf_key))
+    jax.block_until_ready(
+        jitted_filter(pf, filter_model_inputs, pf_init_state, key=pf_filter_key)
+    )
     _times.append(time.perf_counter() - _t0)
 pf_time = float(jnp.median(jnp.array(_times)))
 
diff --git a/docs/examples/kalman_tracking.md b/docs/examples/kalman_tracking.md
index cf7bee1d..d3a3dea3 100644
--- a/docs/examples/kalman_tracking.md
+++ b/docs/examples/kalman_tracking.md
@@ -142,7 +142,8 @@ function with our constructed filter object and model inputs:
 
 ```{.python #kalman-run-filter}
 # Run the filter
-filtered_states = filter(filter_obj, model_inputs, parallel=True)
+init_state = filter_obj.init_prepare(model_inputs[0])
+filtered_states = filter(filter_obj, model_inputs[1:], init_state, parallel=True)
 
 # Extract results
 means = filtered_states.mean  # Posterior state means
@@ -165,7 +166,9 @@ log_normalizing_constant = filtered_states.log_normalizing_constant  # Log margi
 
     ```{.python #kalman-jit-example}
     jitted_filter = jit(filter, static_argnames=['filter_obj', 'parallel'])
-    filtered_states = jitted_filter(filter_obj, model_inputs, parallel=True)
+    filtered_states = jitted_filter(
+        filter_obj, model_inputs[1:], init_state, parallel=True
+    )
     ```
 
 ### Visualize the results
@@ -213,7 +216,10 @@ car's position.
     )
 
     # Run filtering - cuthbert handles NaNs automatically
-    filtered_states_missing = filter(filter_obj_missing, model_inputs_missing)
+    init_state_missing = filter_obj_missing.init_prepare(model_inputs_missing[0])
+    filtered_states_missing = filter(
+        filter_obj_missing, model_inputs_missing[1:], init_state_missing
+    )
 
     print("Successfully handled missing observations!")
     print("Missing observation period: steps 20-30 during the car's turn")
diff --git a/docs/examples/online_stoch_vol.md b/docs/examples/online_stoch_vol.md
index e520df63..d2aa9e63 100644
--- a/docs/examples/online_stoch_vol.md
+++ b/docs/examples/online_stoch_vol.md
@@ -211,8 +211,11 @@ We'll then use the `cuthbert.filter` offline filtering function to run on the pr
 data.
 
 ```{.python #online-stoch-vol-particle-filter-run-previous}
-key, previous_key = random.split(random.key(0))
-previous_states = filter(pf, previous_data, key=key)
+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)
+previous_states = filter(pf, filter_data, init_state, key=previous_key)
 filter_state = tree.map(lambda x: x[-1], previous_states)
 ```
 
@@ -350,4 +353,4 @@ This time the filter sees the actual log-return rather than `jnp.nan`.
 <<online-stoch-vol-particle-filter-filter>>
 <<online-stoch-vol-particle-filter-filter-plot>>
 ```
--->
\ No newline at end of file
+-->
diff --git a/docs/examples/parameter_estimation_em.md b/docs/examples/parameter_estimation_em.md
index 6baf8f44..5b2040b5 100644
--- a/docs/examples/parameter_estimation_em.md
+++ b/docs/examples/parameter_estimation_em.md
@@ -298,7 +298,8 @@ n_epochs = 30
 
 for epoch in range(n_epochs):
     filter_obj, smoother_obj = model_factory(params)
-    filtered_states = filter(filter_obj, model_inputs)
+    init_state = filter_obj.init_prepare(model_inputs[0])
+    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 8c5e25c1..556b7772 100644
--- a/docs/examples/temporal_parallelization_kalman.md
+++ b/docs/examples/temporal_parallelization_kalman.md
@@ -67,15 +67,17 @@ 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])
+filter_model_inputs = model_inputs[1:]
 
 seq_compile_time = timeit.Timer(
     lambda: jax.block_until_ready(
-        jitted_filter(filter_obj, model_inputs, parallel=False)
+        jitted_filter(filter_obj, filter_model_inputs, init_state, parallel=False)
     )
 ).timeit(number=1)
 par_compile_time = timeit.Timer(
     lambda: jax.block_until_ready(
-        jitted_filter(filter_obj, model_inputs, parallel=True)
+        jitted_filter(filter_obj, filter_model_inputs, init_state, parallel=True)
     )
 ).timeit(number=1)
 ```
@@ -88,12 +90,12 @@ num_runs = 10
 
 seq_runtimes = timeit.Timer(
     lambda: jax.block_until_ready(
-        jitted_filter(filter_obj, model_inputs, parallel=False)
+        jitted_filter(filter_obj, filter_model_inputs, init_state, parallel=False)
     )
 ).repeat(repeat=num_runs, number=1)
 par_runtimes = timeit.Timer(
     lambda: jax.block_until_ready(
-        jitted_filter(filter_obj, model_inputs, parallel=True)
+        jitted_filter(filter_obj, filter_model_inputs, init_state, parallel=True)
     )
 ).repeat(repeat=num_runs, number=1)
 
diff --git a/docs/quickstart.md b/docs/quickstart.md
index a48b3307..ff7d726f 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -11,7 +11,7 @@ from typing import NamedTuple
 
 import matplotlib.pyplot as plt
 import pandas as pd
-from jax import Array, vmap
+from jax import Array, tree, vmap
 from jax import numpy as jnp
 from jax.nn import sigmoid
 from jax.scipy.stats import norm
@@ -266,7 +266,10 @@ football_filter = taylor.build_filter(
 We'll use [`cuthbert.filter`][cuthbert.filtering.filter] to easily run offline filtering on our data.
 
 ```{.python #quickstart-run-filter}
-filter_states = filter(football_filter, match_data)
+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)
+filter_states = filter(football_filter, filter_match_data, init_state)
 ```
 
 That was easy wasn't it?
diff --git a/tests/cuthbert/discrete/test_discrete.py b/tests/cuthbert/discrete/test_discrete.py
index a5b25994..fad2dbb3 100644
--- a/tests/cuthbert/discrete/test_discrete.py
+++ b/tests/cuthbert/discrete/test_discrete.py
@@ -133,9 +133,10 @@ def test(self, seed, num_states, num_time_steps, parallel):
         )
 
         # Run the filter and smoother
+        init_state = filter_obj.init_prepare(model_inputs[0])
         filtered_states = self.variant(
             filter, static_argnames=("filter_obj", "parallel")
-        )(filter_obj, model_inputs, parallel=parallel)
+        )(filter_obj, model_inputs[1:], init_state, parallel=parallel)
         filt_dists, log_normalizing_constants = (
             filtered_states.dist,
             filtered_states.log_normalizing_constant,
diff --git a/tests/cuthbert/enkf/test_enkf.py b/tests/cuthbert/enkf/test_enkf.py
index 0e47cff2..b587ec9a 100644
--- a/tests/cuthbert/enkf/test_enkf.py
+++ b/tests/cuthbert/enkf/test_enkf.py
@@ -80,9 +80,10 @@ def test(self, seed, x_dim, y_dim, num_time_steps):
         inference, model_inputs = load_enkf_inference(
             m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys
         )
-        key = random.key(seed + 1)
+        init_key, filter_key = random.split(random.key(seed + 1))
+        init_state = inference.init_prepare(model_inputs[0], key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
-            inference, model_inputs, parallel=False, key=key
+            inference, model_inputs[1:], init_state, parallel=False, key=filter_key
         )
         means = states.mean
         chol_covs = states.chol_cov
@@ -136,10 +137,11 @@ def get_observations(model_inputs):
         )
 
         model_inputs = jnp.arange(num_time_steps + 1)
-        key = random.key(seed + 1)
 
+        init_key, filter_key = random.split(random.key(seed + 1))
+        init_state = inference.init_prepare(model_inputs[0], key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
-            inference, model_inputs, parallel=False, key=key
+            inference, model_inputs[1:], init_state, parallel=False, key=filter_key
         )
 
         # Check shapes
@@ -158,7 +160,15 @@ def init_sample_(key, model_inputs):
                 get_observations=get_observations,
                 n_particles=1_000,
             )
-            states = filter(inference_, model_inputs, parallel=False, key=key)
+            init_key, filter_key = random.split(random.key(seed + 1))
+            init_state = inference_.init_prepare(model_inputs[0], key=init_key)
+            states = filter(
+                inference_,
+                model_inputs[1:],
+                init_state,
+                parallel=False,
+                key=filter_key,
+            )
             return states.log_normalizing_constant[-1]
 
         grad_val = jax.grad(log_nc)(m0)
diff --git a/tests/cuthbert/factorial/test_discrete.py b/tests/cuthbert/factorial/test_discrete.py
index 3227cd5f..c5f850d4 100644
--- a/tests/cuthbert/factorial/test_discrete.py
+++ b/tests/cuthbert/factorial/test_discrete.py
@@ -160,8 +160,9 @@ def test_factorial_discrete_filter(
     )
 
     # output_factorial=False
-    init_state, local_filter_states, final_state = factorial.filter(
-        filter_obj, factorializer, model_inputs, output_factorial=False
+    init_state = filter_obj.init_prepare(model_inputs[0])
+    local_filter_states, final_state = factorial.filter(
+        filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=False
     )
     chex.assert_trees_all_close(init_state.dist, model_params[0], rtol=1e-10, atol=0.0)
     chex.assert_trees_all_close(
@@ -180,7 +181,7 @@ def test_factorial_discrete_filter(
 
     # output_factorial=True
     factorial_states = factorial.filter(
-        filter_obj, factorializer, model_inputs, output_factorial=True
+        filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=True
     )
     chex.assert_trees_all_close(
         factorial_states.dist, factorial_dists_ref, rtol=1e-10, atol=0.0
diff --git a/tests/cuthbert/factorial/test_kalman.py b/tests/cuthbert/factorial/test_kalman.py
index 428ca92a..b0b6823c 100644
--- a/tests/cuthbert/factorial/test_kalman.py
+++ b/tests/cuthbert/factorial/test_kalman.py
@@ -210,8 +210,9 @@ 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, local_filter_states, final_state = factorial.filter(
-        filter_obj, factorializer, model_inputs, output_factorial=False
+    init_state = filter_obj.init_prepare(model_inputs[0])
+    local_filter_states, final_state = factorial.filter(
+        filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=False
     )
     local_filter_covs = (
         local_filter_states.chol_cov
@@ -241,7 +242,7 @@ def test_filter(seed, x_dim, y_dim, num_factors, num_factors_local, num_time_ste
 
     # Check output_factorial = True
     factorial_filtering_states = factorial.filter(
-        filter_obj, factorializer, model_inputs, output_factorial=True
+        filter_obj, factorializer, model_inputs[1:], init_state, output_factorial=True
     )
 
     factorial_filtering_states = cast(ArrayTree, factorial_filtering_states)
@@ -297,8 +298,13 @@ def test_smoother(
         smoother_factorial_index=smoother_factorial_index,
     )
 
-    init_state, local_filter_states, _ = factorial.filter(
-        filter_obj, factorializer, filter_model_inputs, output_factorial=False
+    init_state = filter_obj.init_prepare(filter_model_inputs[0])
+    local_filter_states, _ = factorial.filter(
+        filter_obj,
+        factorializer,
+        filter_model_inputs[1:],
+        init_state,
+        output_factorial=False,
     )
 
     factorial_inds = model_params[-1]
diff --git a/tests/cuthbert/factorial/test_smc.py b/tests/cuthbert/factorial/test_smc.py
index 85a87766..2c434644 100644
--- a/tests/cuthbert/factorial/test_smc.py
+++ b/tests/cuthbert/factorial/test_smc.py
@@ -1,17 +1,16 @@
-from typing import NamedTuple, cast
+from typing import NamedTuple
 
 import chex
 import jax
 import jax.numpy as jnp
 import pytest
-from jax import Array, random
+from jax import Array, random, tree
 
 from cuthbert import factorial
 from cuthbert.inference import Filter
 from cuthbert.smc.particle_filter import ParticleFilterState, build_filter
 from cuthbertlib.resampling import no_resampling, systematic
 from cuthbertlib.stats.multivariate_normal import logpdf
-from cuthbertlib.types import ArrayTree
 from tests.cuthbert.factorial.gaussian_utils import generate_factorial_kalman_model
 from tests.cuthbert.factorial.test_kalman import build_pairwise_factorial_filter
 
@@ -81,7 +80,7 @@ def weighted_mean_and_var(states: ParticleFilterState) -> tuple[Array, Array]:
     return mean, var
 
 
-params = [(0, 8, 2, 8, 3000), (1, 8, 2, 8, 3000)]
+params = [(0, 8, 2, 8, 10000), (1, 8, 2, 8, 10000)]
 
 
 def test_factorial_smc_scalar_particle_leaf():
@@ -130,39 +129,52 @@ def test_factorial_smc_filter(
     kalman_filter, kalman_factorializer, kalman_model_inputs = (
         build_pairwise_factorial_filter(model_params)
     )
-    kalman_states = cast(
-        ArrayTree,
-        factorial.filter(
-            kalman_filter,
-            kalman_factorializer,
-            kalman_model_inputs,
-            output_factorial=True,
-        ),
+    init_state = kalman_filter.init_prepare(kalman_model_inputs[0])
+    init_state = kalman_factorializer.factorialize_init_state(
+        init_state, kalman_model_inputs[0]
+    )
+    kalman_states = factorial.filter(
+        kalman_filter,
+        kalman_factorializer,
+        kalman_model_inputs[1:],
+        init_state,
+        output_factorial=True,
     )
-    _, kalman_local_states, _ = factorial.filter(
-        kalman_filter, kalman_factorializer, kalman_model_inputs, output_factorial=False
+    kalman_local_states, _ = factorial.filter(
+        kalman_filter,
+        kalman_factorializer,
+        kalman_model_inputs[1:],
+        init_state,
+        output_factorial=False,
     )
     kalman_covs = kalman_states.chol_cov @ kalman_states.chol_cov.transpose(0, 1, 3, 2)
 
     smc_filter, smc_factorializer, smc_model_inputs = build_factorial_smc_filter(
         model_params, n_particles=num_particles
     )
-    smc_states = cast(
-        ParticleFilterState,
-        factorial.filter(
-            smc_filter,
-            smc_factorializer,
-            smc_model_inputs,
-            output_factorial=True,
-            key=random.key(seed + 123),
-        ),
+    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
+    )
+    smc_states = factorial.filter(
+        smc_filter,
+        smc_factorializer,
+        smc_filter_model_inputs,
+        init_smc_state,
+        output_factorial=True,
+        key=filter_key,
     )
-    _, smc_local_states, _ = factorial.filter(
+    smc_local_states, _ = factorial.filter(
         smc_filter,
         smc_factorializer,
-        smc_model_inputs,
+        smc_filter_model_inputs,
+        init_smc_state,
         output_factorial=False,
-        key=random.key(seed + 123),
+        key=filter_key,
     )
     kalman_local_covs = (
         kalman_local_states.chol_cov
diff --git a/tests/cuthbert/factorial/test_taylor.py b/tests/cuthbert/factorial/test_taylor.py
index 9f975feb..f160ee05 100644
--- a/tests/cuthbert/factorial/test_taylor.py
+++ b/tests/cuthbert/factorial/test_taylor.py
@@ -60,12 +60,21 @@ def observation_log_density(x, y):
     factorializer = factorial.gaussian.build_factorializer(
         lambda model_inputs: factorial_indices[model_inputs - 1]
     )
-    filter_model_inputs = jnp.arange(num_time_steps + 1)
+    init_model_inputs = jnp.array(0)
+    filter_model_inputs = jnp.arange(1, num_time_steps + 1)
 
-    init_state, local_filter_states, _ = jax.jit(
+    init_state = filter_obj.init_prepare(init_model_inputs)
+    init_state = factorializer.factorialize_init_state(init_state, init_model_inputs)
+    local_filter_states, _ = jax.jit(
         factorial.filter,
         static_argnames=("filter_obj", "factorializer", "output_factorial"),
-    )(filter_obj, factorializer, filter_model_inputs, output_factorial=False)
+    )(
+        filter_obj,
+        factorializer,
+        filter_model_inputs,
+        init_state,
+        output_factorial=False,
+    )
 
     assert init_state.mean.shape == (num_factors, x_dim)
     assert init_state.chol_cov.shape == (num_factors, x_dim, x_dim)
diff --git a/tests/cuthbert/gaussian/test_kalman.py b/tests/cuthbert/gaussian/test_kalman.py
index 8f0b347a..21501598 100644
--- a/tests/cuthbert/gaussian/test_kalman.py
+++ b/tests/cuthbert/gaussian/test_kalman.py
@@ -119,8 +119,10 @@ 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])
+
     # Run sequential sqrt filter
-    seq_states = filter(kalman_filter, model_inputs, parallel=False)
+    seq_states = filter(kalman_filter, model_inputs[1:], init_state, parallel=False)
     seq_means, seq_chol_covs, seq_ells = (
         seq_states.mean,
         seq_states.chol_cov,
@@ -128,7 +130,7 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run parallel sqrt filter
-    par_states = filter(kalman_filter, model_inputs, parallel=True)
+    par_states = filter(kalman_filter, model_inputs[1:], init_state, parallel=True)
     par_means, par_chol_covs, par_ells = (
         par_states.mean,
         par_states.chol_cov,
@@ -166,7 +168,8 @@ def get_loglikelihood_from_params(
         kalman_filter, _, model_inputs = load_kalman_inference(
             m0_, chol_P0_, Fs_, cs_, chol_Qs_, Hs_, ds_, chol_Rs_, ys
         )
-        states = filter(kalman_filter, model_inputs)
+        init_state = kalman_filter.init_prepare(model_inputs[0])
+        states = filter(kalman_filter, model_inputs[1:], init_state)
         return states.log_normalizing_constant[-1]
 
     chex.assert_numerical_grads(
@@ -231,7 +234,8 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    filt_states = filter(kalman_filter, model_inputs)
+    init_state = kalman_filter.init_prepare(model_inputs[0])
+    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)
     Qs = chol_Qs @ chol_Qs.transpose(0, 2, 1)
diff --git a/tests/cuthbert/gaussian/test_moments.py b/tests/cuthbert/gaussian/test_moments.py
index 30b41865..22a2150c 100644
--- a/tests/cuthbert/gaussian/test_moments.py
+++ b/tests/cuthbert/gaussian/test_moments.py
@@ -93,8 +93,10 @@ 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])
+
     # Run sequential sqrt filter
-    seq_states = filter(moments_filter, model_inputs, parallel=False)
+    seq_states = filter(moments_filter, model_inputs[1:], init_state, parallel=False)
     seq_means, seq_chol_covs, seq_ells = (
         seq_states.mean,
         seq_states.chol_cov,
@@ -105,8 +107,15 @@ 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])
+
     # Run associative filter with parallel=False
-    seq_ass_states = filter(associative_moments_filter, model_inputs, parallel=False)
+    seq_ass_states = filter(
+        associative_moments_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=False,
+    )
     seq_ass_means, seq_ass_chol_covs, seq_ass_ells = (
         seq_ass_states.mean,
         seq_ass_states.chol_cov,
@@ -114,7 +123,12 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run associative filter with parallel=True
-    par_ass_states = filter(associative_moments_filter, model_inputs, parallel=True)
+    par_ass_states = filter(
+        associative_moments_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=True,
+    )
     par_ass_means, par_ass_chol_covs, par_ass_ells = (
         par_ass_states.mean,
         par_ass_states.chol_cov,
@@ -200,7 +214,8 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    filt_states = filter(extended_filter, model_inputs)
+    init_state = extended_filter.init_prepare(model_inputs[0])
+    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)
     Qs = chol_Qs @ chol_Qs.transpose(0, 2, 1)
diff --git a/tests/cuthbert/gaussian/test_taylor.py b/tests/cuthbert/gaussian/test_taylor.py
index 2462be0b..7ca06960 100644
--- a/tests/cuthbert/gaussian/test_taylor.py
+++ b/tests/cuthbert/gaussian/test_taylor.py
@@ -153,8 +153,10 @@ 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])
+
     # Run sequential sqrt filter
-    seq_states = filter(taylor_filter, model_inputs, parallel=False)
+    seq_states = filter(taylor_filter, model_inputs[1:], init_state, parallel=False)
     seq_means, seq_chol_covs, seq_ells = (
         seq_states.mean,
         seq_states.chol_cov,
@@ -175,8 +177,15 @@ 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])
+
     # Run associative filter with parallel=False§
-    seq_ass_states = filter(associative_taylor_filter, model_inputs, parallel=False)
+    seq_ass_states = filter(
+        associative_taylor_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=False,
+    )
     seq_ass_means, seq_ass_chol_covs, seq_ass_ells = (
         seq_ass_states.mean,
         seq_ass_states.chol_cov,
@@ -184,7 +193,12 @@ def test_offline_filter(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run associative filter with parallel=True
-    par_ass_states = filter(associative_taylor_filter, model_inputs, parallel=True)
+    par_ass_states = filter(
+        associative_taylor_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=True,
+    )
     par_ass_means, par_ass_chol_covs, par_ass_ells = (
         par_ass_states.mean,
         par_ass_states.chol_cov,
@@ -314,7 +328,8 @@ def test_smoother(seed, x_dim, y_dim, num_time_steps):
     )
 
     # Run the Kalman filter and the standard Kalman smoother.
-    filt_states = filter(log_density_filter, model_inputs)
+    init_state = log_density_filter.init_prepare(model_inputs[0])
+    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)
     Qs = chol_Qs @ chol_Qs.transpose(0, 2, 1)
@@ -421,8 +436,10 @@ 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])
+
     # Run sequential sqrt filter
-    seq_states = filter(taylor_filter, model_inputs, parallel=False)
+    seq_states = filter(taylor_filter, model_inputs[1:], init_state, parallel=False)
     seq_means, seq_chol_covs, seq_ells = (
         seq_states.mean,
         seq_states.chol_cov,
@@ -433,15 +450,27 @@ 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])
+
     # Run associative filter with parallel=False
-    seq_ass_states = filter(associative_taylor_filter, model_inputs, parallel=False)
+    seq_ass_states = filter(
+        associative_taylor_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=False,
+    )
     seq_ass_means, seq_ass_chol_covs, seq_ass_ells = (
         seq_ass_states.mean,
         seq_ass_states.chol_cov,
         seq_ass_states.log_normalizing_constant,
     )
 
-    par_ass_states = filter(associative_taylor_filter, model_inputs, parallel=True)
+    par_ass_states = filter(
+        associative_taylor_filter,
+        model_inputs[1:],
+        associative_init_state,
+        parallel=True,
+    )
     par_ass_means, par_ass_chol_covs, par_ass_ells = (
         par_ass_states.mean,
         par_ass_states.chol_cov,
diff --git a/tests/cuthbert/smc/test_backward_sampler.py b/tests/cuthbert/smc/test_backward_sampler.py
index 226af1f6..fbf9eabc 100644
--- a/tests/cuthbert/smc/test_backward_sampler.py
+++ b/tests/cuthbert/smc/test_backward_sampler.py
@@ -75,8 +75,15 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
         filter_obj, model_inputs, log_potential = load_inference(
             m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys
         )
-        key = random.key(seed + 1)
-        filtered_states = filter(filter_obj, model_inputs, False, key)
+        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)
+        filtered_states = filter(
+            filter_obj,
+            model_inputs[1:],
+            init_state,
+            parallel=False,
+            key=filter_key,
+        )
 
         if method == "tracing":
             bs_fn = tracing
@@ -92,7 +99,6 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
         smoother_obj = build_smoother(
             log_potential, bs_fn, systematic.resampling, n_smoother_particles
         )
-        key, smoother_key = random.split(key)
         smoothed_states = self.variant(
             smoother, static_argnames=("smoother_obj", "parallel")
         )(smoother_obj, filtered_states, model_inputs, False, smoother_key)
@@ -155,10 +161,17 @@ def log_potential(state_prev, state, model_inputs):
             raise ValueError(f"{method} is not a valid backward sampling method.")
 
         key = random.key(0)
-        filter_key, smoother_key = random.split(key)
+        init_key, filter_key, smoother_key = random.split(key, 3)
         num_time_steps = 5
         model_inputs = jnp.empty(num_time_steps + 1)
-        filtered_states = filter(filter_obj, model_inputs, False, filter_key)
+        init_state = filter_obj.init_prepare(model_inputs[0], key=init_key)
+        filtered_states = filter(
+            filter_obj,
+            model_inputs[1:],
+            init_state,
+            parallel=False,
+            key=filter_key,
+        )
 
         n_smoother_particles = 1000
         smoother_obj = build_smoother(
diff --git a/tests/cuthbert/smc/test_particle_filters.py b/tests/cuthbert/smc/test_particle_filters.py
index f2dda632..3de3286a 100644
--- a/tests/cuthbert/smc/test_particle_filters.py
+++ b/tests/cuthbert/smc/test_particle_filters.py
@@ -107,9 +107,14 @@ def test(self, seed, x_dim, y_dim, num_time_steps, method):
         inference, model_inputs = load_inference(
             m0, chol_P0, Fs, cs, chol_Qs, Hs, ds, chol_Rs, ys, method
         )
-        key = random.key(seed + 1)
+        init_key, filter_key = random.split(random.key(seed + 1))
+        init_state = inference.init_prepare(model_inputs[0], key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
-            inference, model_inputs, parallel=False, key=key
+            inference,
+            model_inputs[1:],
+            init_state,
+            parallel=False,
+            key=filter_key,
         )
         weights = jax.nn.softmax(states.log_weights)
         means = jnp.sum(states.particles * weights[..., None], axis=1)
@@ -195,11 +200,17 @@ def log_potential(state_prev, state, model_inputs):
         key = random.key(0)
         num_time_steps = 5
 
-        # Run the particle filter
         model_inputs = jnp.empty(num_time_steps + 1)
-        key, subkey = random.split(key)
+
+        # Run the particle filter
+        init_key, filter_key = random.split(key)
+        init_state = inference.init_prepare(model_inputs[0], key=init_key)
         states = self.variant(filter, static_argnames=("filter_obj", "parallel"))(
-            inference, model_inputs, parallel=False, key=subkey
+            inference,
+            model_inputs[1:],
+            init_state,
+            parallel=False,
+            key=filter_key,
         )
 
         # Verify that the pytree structure is preserved