implement functional init/apply and bind.

docs/flax.linen.rst

@@ -12,7 +12,15 @@ Module
 ------------------------
 
 .. autoclass:: Module
-   :members: setup, variable, param, apply, init, init_with_output, make_rng, variables, Variable, __setattr__
+   :members: setup, variable, param, bind, apply, init, init_with_output, make_rng, variables, Variable, __setattr__
+
+Init/Apply
+------------------------
+
+.. currentmodule:: flax.linen
+.. autofunction:: apply
+.. autofunction:: init
+.. autofunction:: init_with_output
 
 Variables
 ----------------------

examples/vae/train.py

@@ -143,16 +143,17 @@ def loss_fn(params):
 
 @jax.jit
 def eval(params, images, z, z_rng):
-  recon_images, mean, logvar = model().apply({'params': params}, images, z_rng)
+  def eval_model(vae):
+    recon_images, mean, logvar = vae(images, z_rng)
+    comparison = jnp.concatenate([images[:8].reshape(-1, 28, 28, 1),
+                                  recon_images[:8].reshape(-1, 28, 28, 1)])
 
-  comparison = jnp.concatenate([images[:8].reshape(-1, 28, 28, 1),
-                                recon_images[:8].reshape(-1, 28, 28, 1)])
+    generate_images = vae.generate(z)
+    generate_images = generate_images.reshape(-1, 28, 28, 1)
+    metrics = compute_metrics(recon_images, images, mean, logvar)
+    return metrics, comparison, generate_images
 
-  generate_images = model().apply({'params': params}, z, method=VAE.generate)
-  generate_images = generate_images.reshape(-1, 28, 28, 1)
-  metrics = compute_metrics(recon_images, images, mean, logvar)
-
-  return metrics, comparison, generate_images
+  return nn.apply(eval_model, model())({'params': params})
 
 
 def prepare_image(x):

flax/core/__init__.py

@@ -15,5 +15,5 @@
 from .axes_scan import broadcast
 from .frozen_dict import FrozenDict, freeze, unfreeze
 from .tracers import current_trace, trace_level, check_trace_level
-from .scope import Scope, Array, apply, init
+from .scope import Scope, Array, apply, init, bind
 from .lift import scan, vmap, jit

flax/core/scope.py

@@ -290,7 +290,14 @@ def invalidate(self):
     """Invalidates the Scope."""
     self._invalid = True
 
-  def variables(self) -> Collection:
+  def mutable_variables(self) -> VariableDict:
+    """Returns an immutable copy of the mutable variables belonging to this Scope."""
+    self._populate_collections()
+    xs = {k: v for k, v in self._variables.items()
+          if in_filter(self.mutable, k)}
+    return freeze(xs)
+
+  def variables(self) -> VariableDict:
     """Returns an immutable copy of the variables belonging to this Scope."""
     self._populate_collections()
     return freeze(self._variables)
@@ -576,6 +583,29 @@ def _unfreeze_variables(variables, mutable):
   return new_variables
 
 
+def bind(variables: VariableDict,
+         rngs: Optional[RNGSequences] = None,
+         mutable: CollectionFilter = False):
+  """Bind variables and rngs to a new ``Scope``.
+
+  bind provides a ``Scope`` instance without transforming a function
+  with ``apply``. This is particulary useful for debugging and
+  interactive use cases like notebooks where a function would limit
+  the ability split up code into different cells.
+
+  a ``Scope`` instance is a stateful object. Note that idiomatic JAX is functional
+  and therefore a ``Scope` does not mix well well with vanilla JAX APIs. Therefore,
+  we recommend using ``apply`` when code should be reusable and compatible
+  across the JAX software ecosystem.
+  """
+  if not _is_valid_variables(variables):
+    raise errors.ApplyScopeInvalidVariablesError()
+  if rngs is not None and not _is_valid_rngs(rngs):
+    raise errors.ApplyScopeInvalidRngsError()
+  new_variables = _unfreeze_variables(variables, mutable)
+  return Scope(new_variables, rngs=rngs, mutable=mutable)
+
+
 def apply(fn: Callable[..., Any],
           mutable: CollectionFilter = False) -> Callable[..., Any]:
   """Functionalize a `Scope` function.
@@ -593,19 +623,10 @@ def wrapper(variables: VariableDict,
               *args,
               rngs: Optional[RNGSequences] = None,
               **kwargs) -> Union[Any, Tuple[Any, VariableDict]]:
-
-    if not _is_valid_variables(variables):
-      raise errors.ApplyScopeInvalidVariablesError()
-    if rngs is not None and not _is_valid_rngs(rngs):
-      raise errors.ApplyScopeInvalidRngsError()
-    new_variables = _unfreeze_variables(variables, mutable)
-    with Scope(new_variables, rngs=rngs, mutable=mutable).temporary() as root:
+    with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
       y = fn(root, *args, **kwargs)
     if mutable is not False:
-      mutated_variables = {k: v
-                           for k, v in new_variables.items()
-                           if in_filter(mutable, k)}
-      return y, freeze(mutated_variables)
+      return y, root.mutable_variables()
     else:
       return y
 

flax/linen/__init__.py

@@ -25,7 +25,7 @@
                         make_causal_mask, combine_masks)
 from ..core import broadcast
 from .linear import Conv, ConvTranspose, Dense, DenseGeneral, Embed
-from .module import Module, compact, enable_named_call, disable_named_call, Variable
+from .module import Module, compact, enable_named_call, disable_named_call, Variable, init, init_with_output, apply
 from .normalization import BatchNorm, GroupNorm, LayerNorm
 from .pooling import avg_pool, max_pool
 from .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell