def outer_fn(x):
assert x.ndim == 2
x = Bias()(x)
inner = transform.without_apply_rng(transform.transform(inner_fn))
inner_p = lift.lift(inner.init)(base.next_rng_key(), x[0])
vmap_inner = jax.vmap(inner.apply, in_axes=(None, 0))
return vmap_inner(inner_p, x)
def outer_fn(x):
assert x.ndim == 2
x = Bias()(x)
inner = base.transform(inner_fn, state=True)
inner_p, inner_s = lift.lift(inner.init)(base.next_rng_key(), x[0])
vmap_inner = jax.vmap(inner.apply, in_axes=(None, None, 0))
return vmap_inner(inner_p, inner_s, x)[0]
def __call__(self, carry, x):
x += base.get_parameter("w", shape=[], init=jnp.zeros)
inner = transform.transform(inner_fn)
keys = base.next_rng_key() if transform.running_init(
) else None
params = lift.lift(inner.init,
allow_reuse=self._allow_reuse)(keys, x)
return carry, inner.apply(params, None, x)
def test_lift_raises_with_state(self):
f = transform.transform_with_state(
lambda: base.get_state("w", [], init=jnp.zeros))
lifted = lift.lift(f.init) # pytype: disable=wrong-arg-types
with self.assertRaisesRegex(ValueError, "use.*lift_with_state"):
lifted(None)
def test_empty_lift(self):
f = transform.transform(lambda: None)
self.assertEmpty(lift.lift(f.init)(None))