[DanWaxman]

Resolves #281. This fixes an indexing bug in the particle smoothers, which results in the incorrect smoothing distribution. Concretely, state_2.model_inputs was previously used as part of the log_potential:

cuthbert/cuthbert/smc/backward_sampler.py

Lines 181 to 188 in ec53db2

	new_particles_1, ancestors_1 = backward_sampling_fn(
	state_1.key,
	x0_all=state_1.particles,
	x1_all=state_2.particles,
	log_weight_x0_all=state_1.log_weights,
	log_density=lambda s1, s2: log_potential(s1, s2, state_2.model_inputs),
	x1_ancestor_indices=state_2.ancestor_indices,
	)

But the model_inputs convention for particle smoothers differ from the rest of cuthbert, and state_1.model_inputs should be used accordingly:

cuthbert/cuthbert/smc/backward_sampler.py

Lines 128 to 137 in ec53db2

	def smoother_prepare(
	filter_state: ParticleFilterState,
	model_inputs: ArrayTreeLike,
	key: KeyArray | None = None,
	) -> ParticleSmootherState:
	"""Prepare a state for a particle smoother step.
	Note that the model_inputs here are different to filter_state.model_inputs.
	The model_inputs required here are for the transition from t to t+1.
	filter_state.model_inputs represents the transition from t-1 to t.

The proper state_1.model_inputs is now properly used.

This issue wasn't found by the tests because the model definition was wrong in tests: in particular, backward sampling smoothers require the joint log potential, but only the observation log potential was given. We now use the proper model definition, which would fail before but passes now.

Finally, there was a key reshape that caused issues with typed keys. We now split with the correct shape directly, which makes this compatible with both JAX key types.

[Copilot]

Copilot review overview

🟡 Changes recommended

An unused variable fails Ruff checks, and the increased particle count substantially expands quadratic test work.

Review effort: Balanced
Findings: 1 · 1

Open (2)

• Remove unused smoothed_covs assignment · New
• Avoid multiplying expensive smoother test workload · New

What changed in this PR

Fixes particle-smoother input indexing and typed JAX key compatibility.

Changes:

• Uses the correct transition inputs during backward sampling.
• Corrects the smoothing test model.
• Splits MCMC keys directly into the required shape.

File	Description
cuthbert/​smc/​backward_sampler.py	Corrects smoother input indexing.
cuthbertlib/​smc/​smoothing/​mcmc.py	Supports typed and legacy keys.
tests/​cuthbert/​smc/​test_backward_sampler.py	Updates the smoothing regression test.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[SamDuffield]

Good catch and thanks a lot for the fix!!