using flax sow inside a function with scanning

[TalKachman]

Hi Everyone!

Is there a way to use the sow method within a scan function? ideally I would love to save some metadata and the ability to look at states.

Small example of what I want to try:

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

class LSTMSOW(nn.LSTMCell):

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

class LSTMSOW(nn.LSTMCell):
    def __init__(self,*args,**kwargs):
        super.__init(*args,**kwargs)
    def __call__(self,carry,inputs):
        new_carry,outputs = super().__call__(carry,inputs)
        self.sow("intermediates",'output',outputs)

class simpleScan(nn.Module):
    features: int

    @nn.compact
    def __call__(self,xs):
        lstm = nn.scan(LSTMSOW,in_axes=1,out_axes=1,varible_broadcast='params',split_rngs={'params':False})
        dummy_rng = jax.random.PRNGKey(0)
        init_carry = (jnp.zeros((1,self.features)),jnp.zeros((1,self.features)))
        return lstm(init_carry,xs)

key1,key2 = jax.random.split(jax.random.PRNGKey(0))
xs = jax.random.uniform(key1,(1,5,2))
model = simpleScan(2)
init_variables = model.init(key2,xs)
y,state = model.apply(init_variables,xs,mutable=['intermediates'])
print(state)

this gives an empty state. how for example can I access the output within the scan ?
Many thanks!

[andsteing]

You'll need to specify nn.scan(variable_axes={'intermediates': ...} to make nn.Module.sow('intermediates', ...) work inside nn.scan().

For example:

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


class Block(nn.Module):
  width: int
  @nn.compact
  def __call__(self, carry, unused_inputs):
    carry = nn.Dense(self.width)(carry)
    carry = nn.relu(carry)
    self.sow('intermediates', 'carry', carry)
    return carry, None


class MLP(nn.Module):
  width: int
  depth: int
  @nn.compact
  def __call__(self, x):
    self.sow('intermediates', 'x', x)
    carry, unused_outputs = nn.scan(
        Block,
        # If 'intermediates' is not listed in `variable_axes` below, then
        # `self.sow('intermediates', ...)` will not work inside the `nn.scan()`.
        variable_axes={'params': 0, 'intermediates': 0},
        split_rngs={'params': True},
        length=self.depth,
    )(
        width=self.width,
    )(
        x, None,
    )
    return carry


model = MLP(width=2, depth=3)
x = jnp.zeros([1, 2])
variables = model.init(jax.random.PRNGKey(0), x)
out, state = model.apply(variables, x, mutable=['intermediates'])
jax.tree.map(lambda x: x.shape, dict(out=out, state=state, variables=variables))

will output

{'out': (1, 2),
 'state': {'intermediates': {'ScanBlock_0': {'carry': ((3, 1, 2),)},
   'x': ((1, 2),)}},
 'variables': {'params': {'ScanBlock_0': {'Dense_0': {'bias': (3, 2),
     'kernel': (3, 2, 2)}}}}}