capture_module_outputs with Linen API

[rolandgvc]

Hi guys, I was wondering if anyone could provide an example for capturing the intermediate layer outputs when using the Linen API. Does nn.capture_module_outputs still work with Linen? Thanks!

[andsteing]

Note that you can "capture" any intermediate values by adding them to a mutable collection:

from typing import Sequence

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp

class Model(nn.Module):
  features: Sequence[int]

  @nn.compact
  def __call__(self, inputs):
    x = inputs
    for i, features in enumerate(self.features):
      x = nn.relu(nn.Dense(features)(x))
      self.variable('debug', f'dense_relu_{i}', lambda: x).value = x
    return x

model = Model([2, 3])
inputs = jnp.ones([2, 3])
variables = model.init(jax.random.PRNGKey(0), inputs)

model.apply(variables, inputs, mutable=['debug'])

Would output

(DeviceArray([[1.5693309, 0.       , 0.3158317],
              [1.5693309, 0.       , 0.3158317]], dtype=float32), FrozenDict({
     debug: {
         dense_relu_0: DeviceArray([[0.4885028, 1.3949869],
                      [0.4885028, 1.3949869]], dtype=float32),
         dense_relu_1: DeviceArray([[1.5693309, 0.       , 0.3158317],
                      [1.5693309, 0.       , 0.3158317]], dtype=float32),
     },
 }))

[andsteing]

Finally, for reverse-mode wrapping one can use the epsilon trick:

class GradWrapper(nn.Module):
  mdl: Module

  @nn.compact
  def __call__(self, *args, **kwargs):
    y = self.mdl(*args, **kwargs)
    eps = self.variable('inter_grads', 'activation', nn.initializers.zeros, y.shape)
    return y + eps

variables = model.init(...)

grads = jax.grad(model.apply)(variables, batch)
param_grads = grads['params']
inter_grads = grads['inter_grads']

[andsteing]

Please check out https://flax.readthedocs.io/en/latest/howtos/extracting_intermediates.html for different Linen patterns to extract intermediate values.