Merge pull request #3158 from google:haiku-migration-guide-transforms

PiperOrigin-RevId: 542999636

docs/guides/haiku_migration_guide.rst

@@ -571,3 +571,144 @@ explicitly passed, even though the module does not use any stochastic
 operations during ``apply``. In Flax this is not necessary. The Haiku ``rng``
 is set to ``None`` here, but you could also use ``hk.without_apply_rng`` on the
 ``apply`` function to remove the ``rng`` argument.
+
+
+Lifted Transforms
+-----------------
+
+Both Flax and Haiku provide a set of transforms, which we will refer to as lifted transforms,
+that wrap JAX transformations in such a way that they can be used with Modules and sometimes
+provide additional functionality. In this section we will take a look at how to use the
+lifted version of ``scan`` in both Flax and Haiku to implement a simple RNN layer.
+
+To begin, we will first define a ``RNNCell`` module that will contain the logic for a single
+step of the RNN. We will also define a ``initial_state`` method that will be used to initialize
+the state (a.k.a. ``carry``) of the RNN. Like with ``jax.lax.scan``, the ``RNNCell.__call__``
+method will be a function that takes the carry and input, and returns the new
+carry and output. In this case, the carry and the output are the same.
+
+.. codediff::
+  :title_left: Haiku
+  :title_right: Flax
+  :sync:
+
+  class RNNCell(hk.Module):
+    def __init__(self, hidden_size: int, name=None):
+      super().__init__(name=name)
+      self.hidden_size = hidden_size
+
+    def __call__(self, carry, x):
+      x = jnp.concatenate([carry, x], axis=-1)
+      x = hk.Linear(self.hidden_size)(x)
+      x = jax.nn.relu(x)
+      return x, x
+
+    def initial_state(self, batch_size: int):
+      return jnp.zeros((batch_size, self.hidden_size))
+
+  ---
+
+  class RNNCell(nn.Module):
+    hidden_size: int
+
+
+    @nn.compact
+    def __call__(self, carry, x):
+      x = jnp.concatenate([carry, x], axis=-1)
+      x = nn.Dense(self.hidden_size)(x)
+      x = jax.nn.relu(x)
+      return x, x
+
+    def initial_state(self, batch_size: int):
+      return jnp.zeros((batch_size, self.hidden_size))
+
+Next, we will define a ``RNN`` Module that will contain the logic for the entire RNN.
+In Haiku, we will first initialze the ``RNNCell``, then use it to construct the ``carry``,
+and finally use ``hk.scan`` to run the ``RNNCell`` over the input sequence. In Flax its
+done a bit different, we will use ``nn.scan`` to define a new temporary type that wraps
+``RNNCell``. During this process we will also specify instruct ``nn.scan`` to broadcast
+the ``params`` collection (all steps share the same parameters) and to not split the
+``params`` rng stream (so all steps intialize with the same parameters), and finally
+we will specify that we want scan to run over the second axis of the input and stack
+the outputs along the second axis as well. We will then use this temporary type immediately
+to create an instance of the lifted ``RNNCell`` and use it to create the ``carry`` and
+the run the ``__call__`` method which will ``scan`` over the sequence.
+
+.. codediff::
+  :title_left: Haiku
+  :title_right: Flax
+  :sync:
+
+  class RNN(hk.Module):
+    def __init__(self, hidden_size: int, name=None):
+      super().__init__(name=name)
+      self.hidden_size = hidden_size
+
+    def __call__(self, x):
+      cell = RNNCell(self.hidden_size)
+      carry = cell.initial_state(x.shape[0])
+      carry, y = hk.scan(cell, carry, jnp.swapaxes(x, 1, 0))
+      return jnp.swapaxes(y, 0, 1)
+
+  ---
+
+  class RNN(nn.Module):
+    hidden_size: int
+
+    @nn.compact
+    def __call__(self, x):
+      rnn = nn.scan(RNNCell, variable_broadcast='params', split_rngs={'params': False},
+                    in_axes=1, out_axes=1)(self.hidden_size)
+      carry = rnn.initial_state(x.shape[0])
+      carry, y = rnn(carry, x)
+      return y
+
+In general, the main difference between lifted transforms between Flax and Haiku is that
+in Haiku the lifted transforms don't operate over the state, that is, Haiku will handle the
+``params`` and ``state`` in such a way that it keeps the same shape inside and outside of the
+transform. In Flax, the lifted transforms can operate over both variable collections and rng
+streams, the user must define how different collections are treated by each transform
+according to the transforms semantics.
+
+Finally, lets quickly view how the ``RNN`` Module would be used in both Haiku and Flax.
+
+.. codediff::
+  :title_left: Haiku
+  :title_right: Flax
+  :sync:
+
+  def forward(x):
+    return RNN(64)(x)
+
+  model = hk.without_apply_rng(hk.transform(forward))
+
+  params = model.init(
+    PRNGKey(0),
+    x=jax.numpy.ones((3, 12, 32)),
+  )
+
+  y = model.apply(
+    params,
+    x=jax.numpy.ones((3, 12, 32)),
+  )
+
+  ---
+
+  ...
+
+
+  model = RNN(64)
+
+  variables = model.init(
+    PRNGKey(0),
+    x=jax.numpy.ones((3, 12, 32)),
+  )
+  params = variables['params']
+  y = model.apply(
+    {'params': params},
+    x=jax.numpy.ones((3, 12, 32)),
+  )
+
+The only notable change with respect to the examples in the previous sections is that
+this time around we used ``hk.without_apply_rng``in Haiku so that we didn't have to
+the ``rng`` argument as ``None`` to the ``apply``method.