From 3ef8c32c6b67e8d78811229f3be468db8d66a940 Mon Sep 17 00:00:00 2001
From: Jonathan Terhorst <terhorst@gmail.com>
Date: Tue, 21 Oct 2025 16:58:20 -0400
Subject: [PATCH] rematerialize hmm.fwd

---
 src/phlashlib/hmm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/phlashlib/hmm.py b/src/phlashlib/hmm.py
index 4ee5cf9..076ad3f 100644
--- a/src/phlashlib/hmm.py
+++ b/src/phlashlib/hmm.py
@@ -30,6 +30,7 @@ def forward(
 ) -> tuple[Float[Array, "M"], Scalar]:
     emis = jnp.stack([pp.emis0, pp.emis1, jnp.ones_like(pp.emis0)])
 
+    @jax.remat
     def fwd(tup, ob):
         alpha_hat, ll = tup
         alpha_hat = _matvec_smc(alpha_hat, pp)