Fix scan_checkpoints with padded sequences

Author: ricardoV94

pytensor/scan/checkpoints.py

@@ -1,7 +1,7 @@
 import pytensor.tensor.basic as ptb
 from pytensor.scan.basic import scan
 from pytensor.tensor.basic import Join
-from pytensor.tensor.math import ceil, eq
+from pytensor.tensor.math import ceil, eq, neq
 from pytensor.tensor.subtensor import set_subtensor
 
 
@@ -130,16 +130,18 @@ def scan_checkpoints(
         # Since padding could be an empty tensor, Join returns a view of s.
         join = Join(view=0)
         for i, s in enumerate(sequences):
-            n = s.shape[0] % save_every_N
-            z = ptb.zeros((n, s.shape[1:]), dtype=s.dtype)
-            sequences[i] = join(0, [s, z])
+            overshoots_by = s.shape[0] % save_every_N
+            overshoots = neq(overshoots_by, 0)
+            n = (save_every_N - overshoots_by) * overshoots
+            z = ptb.zeros((n, *s.shape[1:]), dtype=s.dtype)
+            sequences[i] = join(0, s, z)
 
     # Establish the input variables of the outer scan
     o_sequences = [
         s.reshape(
-            [s.shape[0] / save_every_N, save_every_N]
+            [s.shape[0] // save_every_N, save_every_N]
             + [s.shape[i] for i in range(1, s.ndim)],
-            s.ndim + 1,
+            ndim=s.ndim + 1,
         )
         for s in sequences
     ]

tests/scan/test_checkpoints.py

@@ -5,23 +5,26 @@
 from pytensor.gradient import grad
 from pytensor.scan.basic import scan
 from pytensor.scan.checkpoints import scan_checkpoints
-from pytensor.tensor.basic import ones_like
+from pytensor.tensor.basic import arange, ones_like
 from pytensor.tensor.type import iscalar, vector
 
 
 class TestScanCheckpoint:
     def setup_method(self):
         self.k = iscalar("k")
         self.A = vector("A")
+        seq = arange(self.k, dtype="float32") + 1
         result, _ = scan(
-            fn=lambda prior_result, A: prior_result * A,
+            fn=lambda s, prior_result, A: prior_result * A / s,
             outputs_info=ones_like(self.A),
+            sequences=[seq],
             non_sequences=self.A,
             n_steps=self.k,
         )
         result_check, _ = scan_checkpoints(
-            fn=lambda prior_result, A: prior_result * A,
+            fn=lambda s, prior_result, A: prior_result * A / s,
             outputs_info=ones_like(self.A),
+            sequences=[seq],
             non_sequences=self.A,
             n_steps=self.k,
             save_every_N=100,