def test_stateful_fori_with_rng_use(self, iteration_count):
tmp_default = base.DEFAULT_PRNG_RESERVE_SIZE
base.DEFAULT_PRNG_RESERVE_SIZE = 64
def body_fun(_, x):
for _ in range(10):
_ = base.next_rng_key()
return x
base.reserve_rng_keys(5)
_ = stateful.fori_loop(0, iteration_count, body_fun, 1)
base.DEFAULT_PRNG_RESERVE_SIZE = tmp_default
def f(lower, upper): y = stateful.fori_loop(lower, upper, lambda i, x: m(i), 2) return y
def test_fori(self, lower, n): upper = lower + n m = CountingModule() y = stateful.fori_loop(lower, upper, lambda i, x: m(i), 2) self.assertEqual(y, jnp.square(upper - 1)) self.assertEqual(m.count, upper - lower)
("named_call", stateful.named_call),
# Parallelization.
# TODO(tomhennigan): Add missing features (e.g. pjit,xmap).
# ("pmap", lambda f: stateful.pmap(f, "i")),
# Vectorization.
("vmap", lambda f: stateful.vmap(f, split_rng=False)),
# Control flow.
# TODO(tomhennigan): Enable for associative_scan.
# ("associative_scan", lambda f:
# (lambda x: jax.lax.associative_scan(f, x))),
("cond", lambda f: (lambda x: stateful.cond(True, f, f, x))),
("fori_loop", lambda f:
(lambda x: stateful.fori_loop(0, 1, base_test.ignore_index(f), x))),
# ("map", lambda f: (lambda x: stateful.map(f, x))),
("scan", lambda f:
(lambda x: stateful.scan(base_test.identity_carry(f), None, x))),
("switch", lambda f: (lambda x: stateful.switch(0, [f, f], x))),
("while_loop", lambda f: toggle(
f, lambda x: stateful.while_loop(lambda xs: xs[0] == 0,
lambda xs: (1, f(xs[1])),
(0, x)))),
# Automatic differentiation.
# TODO(tomhennigan): Add missing features (e.g. custom_vjp, custom_jvp).
("grad", lambda f: stateful.grad(lambda x: f(x).sum())),
("value_and_grad", lambda f: stateful.value_and_grad(lambda x: f(x).sum())),
("checkpoint", stateful.remat),