mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #19634 from jakevdp:key-reuse-scan

Browse Source 
PiperOrigin-RevId: 604418753
This commit is contained in:
jax authors 2024-02-05 13:39:07 -08:00
commit c1c0c1cf33
3 changed files with 25 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
jax/experimental/key_reuse/_forwarding.py
View File

@ -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=}")

14
jax/experimental/key_reuse/_simple.py
View File

@ -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=}")

7
tests/key_reuse_test.py
View File

@ -710,6 +710,13 @@ class KeyReuseIntegrationTest(jtu.JaxTestCase):
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):