def test_without_state(self):
def f():
w = base.get_parameter("w", [], init=jnp.zeros)
return w
init_fn, apply_fn = base.without_state(base.transform_with_state(f))
params = init_fn(None)
out = apply_fn(params, None)
self.assertEqual(out, 0)
def test_without_state_raises_if_state_used(self):
def f():
for _ in range(10):
count = base.get_state("count", (), jnp.int32, jnp.zeros)
base.set_state("count", count + 1)
return count
init_fn, _ = base.without_state(base.transform_with_state(f))
with self.assertRaisesRegex(ValueError, "use.*transform_with_state"):
init_fn(None)
def test_without_state_raises_if_state_used(self):
init_fn, _ = base.without_state(
base.transform_with_state(lambda: CountingModule()())) # pylint: disable=unnecessary-lambda
with self.assertRaisesRegex(ValueError, "use.*transform_with_state"):
init_fn(None)
def test_without_state(self):
init_fn, apply_fn = base.without_state(
base.transform_with_state(lambda: ScalarModule()())) # pylint: disable=unnecessary-lambda
params = init_fn(None)
out = apply_fn(params, None)
self.assertEqual(out, 0)