From 02055301140d1d0afb5b694ecf716fae1329db7e Mon Sep 17 00:00:00 2001 From: saikyo219 Date: Sun, 23 Aug 2026 02:18:51 +0900 Subject: [PATCH] Fix generator being appended instead of a list in scan carry_shapes A missing pair of brackets in the enum scan path made carry_shapes collect a generator object instead of a list of shapes (the sibling append in body_fn builds a list). This currently works only because the generator's iterable is evaluated eagerly and each entry is consumed at most once; a second consumption would silently yield nothing. Making it a list also renders the ty ignore on the sibling append unnecessary, since the element type of carry_shapes is now consistent. Co-Authored-By: Claude Fable 5 --- numpyro/contrib/control_flow/scan.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py index 509e987a4..c7a21b202 100644 --- a/numpyro/contrib/control_flow/scan.py +++ b/numpyro/contrib/control_flow/scan.py @@ -192,7 +192,7 @@ def body_fn(wrapped_carry, x, prefix=None): # store shape of new_carry at a global variable if len(carry_shapes) < (history + 1): carry_shapes.append( - [jnp.shape(x) for x in jax.tree.flatten(new_carry)[0]] # ty: ignore[invalid-argument-type] + [jnp.shape(x) for x in jax.tree.flatten(new_carry)[0]] ) # make new_carry have the same shape as carry # FIXME: is this rigorous? @@ -225,7 +225,7 @@ def body_fn(wrapped_carry, x, prefix=None): # shape so we don't need to record them here if (i >= history - 1) and (len(carry_shapes) < history + 1): carry_shapes.append( - jnp.shape(x) for x in jax.tree.flatten(wrapped_carry[-1])[0] + [jnp.shape(x) for x in jax.tree.flatten(wrapped_carry[-1])[0]] ) else: # this is the last rolling step