Repository: google-deepmind/dm-haiku

Issue:
Is there a way to share parameters between methods?
You can write and transform multiple methods on the same module, but it doesn't seem possible to share parameters between them without manually merging the two parameter FlatMappings. It's particularly cumbersome if the shared parameters are used several submodules deep. Is there any more convenient way to accomplish something like this?


Discussion:
Hi @davisyoshida , there are two options I'm aware of. Either we have a utility to combine the parameters (`hk.data_structures.merge`) or there is a (somewhat common internally) pattern where you can create "multi transform"  a couple of functions together and give them a shared init.

I've put together an example of both here: https://colab.research.google.com/gist/tomhennigan/cae5ffc066bee24eca56e85d8bf4d351/example-of-transforming-multiple-methods.ipynb
@tomhennigan Thanks! The second one seems really useful. Are there edge cases or something which prevent it from being added to Haiku officially?
I don't think there are any fundamental blockers to this becoming part of Haiku, I'll follow up with the folks who originally came up with this pattern next week and see if we can upstream their impl (since it has been battle tested by a large number of researchers).
@tomhennigan Looking forward to this being added! I would also love to see some examples of how this is used by researchers at DM. I am recently working on RL using Soft Actor-Critic. Having to perform transforms of multiple methods and sharing parameters between them quicks become unmanageable!
@tomhennigan - thanks so much for going into detail on this! Super helpful.

Given the two options you describe of (1) combining the parameters (`hk.data_structures.merge`) or (2) the 'multi transform' option, I am finding (1) results in more readable and shorter haiku code:

https://github.com/altosaar/variational-autoencoder/blob/f21223271a6a73c0fd3d40492dff485fd4798cf7/train_variational_autoencoder_jax.py

I tried the 'multi transform' option, but having a class for VariationalAutoEncoder used more code than applying `hk.transform` separately to the model and variational distributions. The added benefit of keeping it separate is flexibility.

The downside to combining the parameters is a lot of boilerplate:

```
        predicate = lambda module_name, name, value: "model" in module_name
        p_params, q_params = hk.data_structures.partition(predicate, params)
        q_z = variational.apply(q_params, rng_key, x)
        z = q_z.sample(sample_shape=[1], seed=rng_key)
        p_z, p_x_given_z = model.apply(p_params, rng_key, x, z)
```

(This replaces the __call__ method of the multi transform class.)

Are there any ways to avoid such boilerplate while avoiding multi transform functions/classes?

Thanks again! Excited to make more code public using this framework.

Resolving patch:
diff --git a/haiku/__init__.py b/haiku/__init__.py
--- a/haiku/__init__.py
+++ b/haiku/__init__.py
@@ -60,6 +60,10 @@
 from haiku._src.module import transparent
 from haiku._src.moving_averages import EMAParamsTree
 from haiku._src.moving_averages import ExponentialMovingAverage
+from haiku._src.multi_transform import multi_transform
+from haiku._src.multi_transform import multi_transform_with_state
+from haiku._src.multi_transform import MultiTransformed
+from haiku._src.multi_transform import MultiTransformedWithState
 from haiku._src.pool import avg_pool
 from haiku._src.pool import AvgPool
 from haiku._src.pool import max_pool
@@ -141,6 +145,8 @@
     "MaxPool",
     "Module",
     "MultiHeadAttention",
+    "MultiTransformed",
+    "MultiTransformedWithState",
     "PRNGSequence",
     "Params",
     "RNNCore",
@@ -171,6 +177,8 @@
     "max_pool",
     "maybe_next_rng_key",
     "mixed_precision",
+    "multi_transform",
+    "multi_transform_with_state",
     "multinomial",
     "nets",
     "next_rng_key",
diff --git a/haiku/_src/multi_transform.py b/haiku/_src/multi_transform.py
new file mode 100644
--- /dev/null
+++ b/haiku/_src/multi_transform.py
@@ -0,0 +1,257 @@
+# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Transform a set of Haiku-using functions which use overlapping params."""
+
+# pylint: disable=unnecessary-lambda
+
+import types
+from typing import Any, Callable, NamedTuple, Tuple
+
+import dataclasses
+from haiku._src import analytics
+from haiku._src import transform
+from haiku._src import typing
+import jax
+
+# If you are forking replace this with `import haiku as hk`.
+hk = types.ModuleType('haiku')
+hk.transform_with_state = transform.transform_with_state
+hk.Params = typing.Params
+hk.State = typing.State
+del transform, typing
+
+PyTreeDef = Any
+TemplateFn = Callable[..., Any]
+TreeOfApplyFns = Any
+
+
+class MultiTransformed(NamedTuple):
+  """Holds a collection of pure functions.
+
+  Attributes:
+    init: A pure function: ``params = init(rng, *a, **k)``
+    apply: A JAX tree of pure functions each with the signature:
+      ``out = apply(params, rng, *a, **k)``.
+
+  See also:
+    :class:`Transformed`: Single apply variant of multi-transform.
+    :class:`MultiTransformedWithState`: Multi apply with state variant.
+  """
+
+  # Args: [Optional[PRNGKey], ...]
+  init: Callable[..., hk.Params]
+
+  # PyTree[Callable[[hk.Params, Optional[PRNGKey], ...], Any]]
+  apply: Any
+
+
+class MultiTransformedWithState(NamedTuple):
+  """Holds a collection of pure functions.
+
+  Attributes:
+    init: A pure function: ``params, state = init(rng, *a, **k)``
+    apply: A JAX tree of pure functions each with the signature:
+      ``out, state = apply(params, state, rng, *a, **k)``.
+
+  See also:
+    :class:`TransformedWithState`: Single apply variant of multi-transform.
+    :class:`MultiTransformed`: Multi apply with state variant.
+  """
+
+  # Args: [Optional[PRNGKey], ...]
+  init: Callable[..., Tuple[hk.Params, hk.State]]
+
+  # PyTree[Callable[[hk.Params, hk.State, Optional[PRNGKey], ...],
+  #                 Tuple[Any, hk.State]]]
+  apply: Any
+
+
+@dataclasses.dataclass
+class Box:
+  """Holds a Python value and has no leaves."""
+  python_value: Any
+
+jax.tree_util.register_pytree_node(
+    Box, lambda b: ([], b.python_value), lambda v, _: Box(v))
+
+
+def multi_transform_with_state(
+    f: Callable[[], Tuple[TemplateFn, TreeOfApplyFns]],
+) -> MultiTransformedWithState:
+  """Transforms a collection of functions using Haiku into pure functions.
+
+  See :func:`multi_transform` for more details.
+
+  Example:
+
+  >>> def f():
+  ...   encoder = hk.Linear(1, name="encoder")
+  ...   decoder = hk.Linear(1, name="decoder")
+  ...
+  ...   def init(x):
+  ...     z = encoder(x)
+  ...     return decoder(z)
+  ...
+  ...   return init, (encoder, decoder)
+
+  >>> f = hk.multi_transform_with_state(f)
+  >>> rng = jax.random.PRNGKey(42)
+  >>> x = jnp.ones([1, 1])
+  >>> params, state = f.init(rng, x)
+  >>> jax.tree_map(jnp.shape, params)
+  {'decoder': {'b': (1,), 'w': (1, 1)},
+   'encoder': {'b': (1,), 'w': (1, 1)}}
+
+  >>> encode, decode = f.apply
+  >>> z, state = encode(params, state, None, x)
+  >>> y, state = decode(params, state, None, z)
+
+  Args:
+    f: Function returning a "template" function and an arbitrary
+      tree of functions using modules connected in the template function.
+
+  Returns:
+    An ``init`` function and a tree of pure ``apply`` functions.
+
+  See also:
+    :func:`transform_with_state`: Transform a single apply function.
+    :func:`multi_transform`: Transform multiple apply functions without state.
+  """
+  analytics.log_once('multi_transform_with_state')
+
+  def init_fn(*args, **kwargs):
+    """Returns initial state for the transformed functions."""
+    return f()[0](*args, **kwargs)
+
+  init_fn = hk.transform_with_state(init_fn).init
+
+  def apply_fn_i(i):
+    def apply_fn(*args, **kwargs):
+      """Applies the transformed function at the given inputs."""
+      return jax.tree_leaves(f()[1])[i](*args, **kwargs)
+    return apply_fn
+
+  # We need to find out the structure of f()[1], including how many
+  # functions there are, so that we can transform them individually and repack
+  # into the same tree structure. It's valid for modules to declare parameters
+  # in their constructor, so we need to create something that looks like
+  # hk.Params in order to do this. `jax.eval_shape` interprets the function
+  # abstractly, ie no real params are created, and we don't need to touch the
+  # accelerator. This means hardcoding the RNG below is fine.
+  def get_output_treedef() -> Box:
+    rng = jax.random.PRNGKey(42)  # This is fine, see above
+    fns = hk.transform_with_state(lambda: f()[1])
+    apply_fns, _ = fns.apply(*fns.init(rng), rng)
+    return Box(jax.tree_structure(apply_fns))
+
+  output_treedef = jax.eval_shape(get_output_treedef).python_value
+  apply_fns = make_tree(lambda i: hk.transform_with_state(apply_fn_i(i)).apply,
+                        output_treedef)
+
+  return MultiTransformedWithState(init_fn, apply_fns)
+
+
+def multi_transform(
+    f: Callable[[], Tuple[TemplateFn, TreeOfApplyFns]],
+) -> MultiTransformed:
+  """Transforms a collection of functions using Haiku into pure functions.
+
+  In many scenarios we have several modules which are used either as primitives
+  for several Haiku modules/functions, or whose pure versions are to be reused
+  in downstream code. This utility enables this by applying
+  :func:`transform` to an arbitrary tree of Haiku functions which share modules
+  and have a common ``init`` function.
+
+  ``f`` is expected to return a tuple of two elements. First is a ``template``
+  Haiku function which provides an example of how all internal Haiku modules are
+  connected. This function is used to create a common ``init`` function (with
+  your parameters).
+
+  The second object is an arbitrary tree of Haiku functions all of which reuse
+  the modules connected in the ``template`` function. These functions are
+  transformed to pure ``apply`` functions.
+
+  Example:
+
+  >>> def f():
+  ...   encoder = hk.Linear(1, name="encoder")
+  ...   decoder = hk.Linear(1, name="decoder")
+  ...
+  ...   def init(x):
+  ...     z = encoder(x)
+  ...     return decoder(z)
+  ...
+  ...   return init, (encoder, decoder)
+
+  >>> f = hk.multi_transform(f)
+  >>> rng = jax.random.PRNGKey(42)
+  >>> x = jnp.ones([1, 1])
+  >>> params = f.init(rng, x)
+  >>> jax.tree_map(jnp.shape, params)
+  {'decoder': {'b': (1,), 'w': (1, 1)},
+   'encoder': {'b': (1,), 'w': (1, 1)}}
+
+  >>> encode, decode = f.apply
+  >>> z = encode(params, None, x)
+  >>> y = decode(params, None, z)
+
+  Args:
+    f: A factory function that returns two functions, firstly a common init
+      function that creates all modules, and secondly a pytree of apply
+      functions which make use of those modules.
+
+  Returns:
+    A :class:`MultiTransformed` instance which contains a pure init function
+      that creates all parameters, and a pytree of pure apply functions that
+      given the params apply the given function.
+
+  See also:
+    :func:`multi_transform_with_state`: Equivalent for modules using state.
+  """
+  analytics.log_once('multi_transform')
+
+  f = multi_transform_with_state(f)
+  f = without_state(f)
+  return f
+
+
+def without_state(f: MultiTransformedWithState) -> MultiTransformed:
+  """Converts ``MultiTransformedWithState`` to ``MultiTransformed``."""
+  def init_fn(rng, *args, **kwargs) -> hk.Params:
+    params, state = f.init(rng, *args, **kwargs)
+    if state:
+      raise ValueError(
+          'If your transformed function uses `hk.{get,set}_state` then use '
+          '`hk.multi_transform_with_state`.')
+    return params
+
+  def apply_without_state(orig_apply_fn) -> Any:
+    def apply_fn(params: hk.Params, rng, *args, **kwargs):
+      out, state = orig_apply_fn(params, {}, rng, *args, **kwargs)
+      if state:
+        raise ValueError(
+            'If your transformed function uses `hk.{get,set}_state` then use '
+            '`hk.multi_transform_with_state`.')
+      return out
+    return apply_fn
+
+  apply_fns = jax.tree_map(apply_without_state, f.apply)
+
+  return MultiTransformed(init_fn, apply_fns)
+
+
+def make_tree(f: Callable[[int], Any], treedef: PyTreeDef):
+  leaves = list(map(f, range(treedef.num_leaves)))
+  return jax.tree_unflatten(treedef, leaves)


Test patch:
diff --git a/haiku/_src/multi_transform_test.py b/haiku/_src/multi_transform_test.py
new file mode 100644
--- /dev/null
+++ b/haiku/_src/multi_transform_test.py
@@ -0,0 +1,108 @@
+# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for haiku._src.multi_transform."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from haiku._src import base
+from haiku._src import multi_transform
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def _assert_arrays_equal(x: jnp.DeviceArray, y: jnp.DeviceArray) -> None:
+  np.testing.assert_almost_equal(np.array(x), np.array(y), decimal=5)
+
+
+class MultiTransformTest(parameterized.TestCase):
+
+  @parameterized.parameters(multi_transform.multi_transform,
+                            multi_transform.multi_transform_with_state)
+  def test_multi_transform_empty(self, mt):
+    for empty_tree in ({}, [], (), AttrMap()):
+      with self.subTest(type(empty_tree).__name__):
+        f = mt(lambda: (lambda: None, empty_tree))  # pylint: disable=cell-var-from-loop
+        f.init(None)
+        self.assertEqual(f.apply, empty_tree)
+
+  def test_custom_pytree(self):
+    def f():
+      init = lambda: None
+      foo = lambda: 'foo'
+      bar = lambda: 'bar'
+      return init, AttrMap(foo=foo, bar=bar)
+
+    f = multi_transform.multi_transform(f)
+    self.assertEqual('foo', f.apply.foo({}, None))
+    self.assertEqual('bar', f.apply.bar({}, None))
+
+  def test_parameter_in_init(self):
+    def f():
+      w = base.get_parameter('w', [], init=jnp.zeros)
+      s = base.get_state('s', [], init=jnp.zeros)
+      init = lambda: None
+
+      def add():
+        s_add = base.get_state('s', [], init=jnp.zeros)
+        w_add = base.get_parameter('w', [], init=jnp.zeros)
+        return w, w_add, s, s_add
+
+      def sub():
+        s_sub = base.get_state('s', [], init=jnp.zeros)
+        w_sub = base.get_parameter('w', [], init=jnp.zeros)
+        return w, w_sub, s, s_sub
+
+      return init, (add, sub)
+
+    f = multi_transform.multi_transform_with_state(f)
+    params, state = f.init(None)
+    self.assertLen(f.apply, 2)
+    for apply_fn in f.apply:
+      # Check parameter and state reuse inside the transformed function.
+      (w, w_apply, s, s_apply), _ = apply_fn(params, state, None)
+      self.assertIs(w, w_apply)
+      self.assertIs(s, s_apply)
+
+  def test_state(self):
+    def f():
+      def init():
+        s = base.get_state('s', [], init=jnp.zeros)
+        base.set_state('s', s + 1)
+
+      def apply():
+        s = base.get_state('s')
+        base.set_state('s', s + 1)
+      return init, apply
+
+    f = multi_transform.multi_transform_with_state(f)
+    _, state_in = f.init(None)
+    self.assertEqual(state_in, {'~': {'s': 0}})
+    _, state_out = f.apply({}, state_in, None)
+    self.assertEqual(state_out, {'~': {'s': 1}})
+
+
+# Example custom pytree (a dict where `x.a` behaves like `x['a']`).
+class AttrMap(dict):
+  __getattr__ = dict.__getitem__
+  __setattr__ = dict.__setitem__
+  __delattr__ = dict.__delitem__
+
+jax.tree_util.register_pytree_node(AttrMap,
+                                   lambda d: (list(d.values()), d.keys()),
+                                   lambda k, v: AttrMap(zip(k, v)))
+
+if __name__ == '__main__':
+  absltest.main()