def f(x):
@lu.wrap_init
@jax.named_scope('bar')
def _f(x):
return [x + 1]
return core.call(_f, x)[0]
def f(x):
@lu.wrap_init
def f_(x):
effect_p.bind(effect='foo')
effect_p.bind(effect='bar')
return [x]
return core.call(f_, x)[0]
def f(x):
@lu.wrap_init
@extend_name_stack('bar')
def _f(x):
return [x + 1]
return core.call(_f, x)[0]