diff --git a/jax/experimental/key_reuse/_forwarding.py b/jax/experimental/key_reuse/_forwarding.py index db5a9dc9359f..3e46c8b12ab8 100644 --- a/jax/experimental/key_reuse/_forwarding.py +++ b/jax/experimental/key_reuse/_forwarding.py @@ -260,16 +260,20 @@ def _scan_key_type_signature(eqn, args_consumed): # scan body should not consume key in constants if any(np.any(s.mask) for s in signature.sinks if s.idx < num_consts): - raise KeyReuseError("scan body function leads to key reuse when repeatedly executed:\n" + raise KeyReuseError("scan body function leads to key reuse when repeatedly executed, " + "because key constants are repeatedly consumed:\n" f" {signature=}\n" f" {eqn=}\n" f" {jaxpr=}") # scan carry should only consume keys that are sourced on output. - carry_sinks = {s.idx - num_consts: s.mask for s in signature.sinks if 0 <= s.idx - num_consts < num_carry} - carry_sources = {s.idx: s.mask for s in signature.sources if 0 <= s.idx < num_carry} - if carry_sinks.keys() != carry_sources.keys(): # TODO(jakevdp): check that masks match - raise KeyReuseError(f"scan body function leads to key reuse when repeatedly executed:\n" + carry_sinks = {s.idx - num_consts: s.mask for s in signature.sinks + if 0 <= s.idx - num_consts < num_carry and np.any(s.mask)} + carry_sources = {s.idx: s.mask for s in signature.sources + if 0 <= s.idx < num_carry and np.any(s.mask)} + if not set(carry_sinks).issubset(set(carry_sources)): # TODO(jakevdp): check that masks match + raise KeyReuseError("scan body function leads to key reuse when repeatedly executed, " + "because consumed inputs don't match sourced outputs:\n" f" {signature=}\n" f" {eqn=}\n" f" {jaxpr=}") diff --git a/jax/experimental/key_reuse/_simple.py b/jax/experimental/key_reuse/_simple.py index edc4f2a460db..0d379587cf65 100644 --- a/jax/experimental/key_reuse/_simple.py +++ b/jax/experimental/key_reuse/_simple.py @@ -231,16 +231,20 @@ def _scan_key_type_signature(eqn, args_consumed): # scan body should not consume key in constants if any(np.any(s.mask) for s in signature.sinks if s.idx < num_consts): - raise KeyReuseError(f"scan body function leads to key reuse when repeatedly executed:\n" + raise KeyReuseError("scan body function leads to key reuse when repeatedly executed, " + "because key constants are repeatedly consumed:\n" f" {signature=}\n" f" {eqn=}\n" f" {jaxpr=}") # scan carry should only consume keys that are sourced on output. - carry_sinks = {s.idx - num_consts: s.mask for s in signature.sinks if 0 <= s.idx - num_consts < num_carry} - carry_sources = {s.idx: s.mask for s in signature.sources if 0 <= s.idx < num_carry} - if carry_sinks.keys() != carry_sources.keys(): # TODO(jakevdp): check that masks match - raise KeyReuseError(f"scan body function leads to key reuse when repeatedly executed:\n" + carry_sinks = {s.idx - num_consts: s.mask for s in signature.sinks + if 0 <= s.idx - num_consts < num_carry and np.any(s.mask)} + carry_sources = {s.idx: s.mask for s in signature.sources + if 0 <= s.idx < num_carry and np.any(s.mask)} + if not set(carry_sinks).issubset(set(carry_sources)): + raise KeyReuseError("scan body function leads to key reuse when repeatedly executed, " + "because consumed inputs don't match sourced outputs:\n" f" {signature=}\n" f" {eqn=}\n" f" {jaxpr=}") diff --git a/tests/key_reuse_test.py b/tests/key_reuse_test.py index 5ebecf4d4e19..0cda84d99b78 100644 --- a/tests/key_reuse_test.py +++ b/tests/key_reuse_test.py @@ -710,6 +710,13 @@ def f_scan_over_keys(key): return jax.lax.map(jax.random.bits, keys) self.check_key_reuse(f_scan_over_keys, jax.random.key(0)) + def test_scan_consume_one(self): + def f_scan_over_keys(*keys): + def body_func(keys, x): + return tuple(jax.random.split(keys[0])), x + return jax.lax.scan(body_func, keys, xs=jnp.arange(10)) + self.check_key_reuse(f_scan_over_keys, jax.random.key(0), jax.random.key(1)) + def test_vmap(self): @jax.vmap def f_good(seed):