How to pass submodule (as pure function) to another JAX library, which might use jax function transform and lax control flow internally?

[YouJiacheng]

Simple example:

def foo(fun, ...):
    return jax.lax.scan(fun, ...)
class Model(nn.Module)
    @nn.compact
    def __call__(self, ...):
        fun = submodule()
        return foo(fun, ...)

Solution

Assume submodule only have params variable

class Model(nn.Module)
    @nn.compact
    def __call__(self, x):
        m = submodule()
        m_params = self.param('m', m.init, x)
        return foo(lambda x: m.apply(m_params, x), ...)

[jheek]

You can use flax.linen.scan for a scan within a Module

[YouJiacheng]

I know, but I want to use jax.experimental.ode or diffrax, which use scan internally...
I could use it outside module, but it will break my abstraction...

[jheek]

If you are using just params and no variables and RNGs than it's still pretty easy to call lax.scan. See https://flax.readthedocs.io/en/latest/design_notes/lift.html for how you can do this

[YouJiacheng]

@jheek Thanks! It seems that this method perfectly solve my problem!
I have read this document before, but I didn't notice that this method actually convert a module into a pure function!
Just need to delegate self.param's init_fn to submodule.init and use submodule.apply instead of __call__!
One more question:
For the first limitation

During initialization, we call the submodule twice through init_fn and apply_fn. If the submodule used the same trick to do functional transformation we will end up executing a lot of code as the number of module calls grows like 2^d where d is the number of nested function transformations.

I think there is something wrong: when we call submodule through apply_fn, submodule won't be initialized again, and will only call sub-submodule once through apply_fn, thus we only call d-nested submodule d times instead of 2^d times.

BTW, all apply_fn called during initialization is only abstract evaluation on shape and dtype(no actual computation happens), right?(Oh, this is true only if I use jit to eliminate useless computation?)
Could I use following method to overcome the first limitation?

class Model(nn.Module):
  @nn.compact
  def __call__(self, x):
    mlp = MLP(parent=None)
    o = None
    def init_fn(rng, x):
        nonlocal o
        o, v = mlp.init_with_output(rng, x)
        return v['params']
    apply_fn = mlp.apply
    mlp_params = self.param('mlp', init_fn, xs)
    if self.is_mutable_collection('params'):
        return o
    return apply_fn({'params': mlp_params}, xs)

I noticed that after initialization, self.param will use jax.eval_shape on init_fn to check the shape, but it seems that the side-effect is not important.

[cgarciae]

@YouJiacheng The pattern is:

class Model(nn.Module)
    @nn.compact
    def __call__(self, x):
        module = submodule()
        def scan_fn(module, carry, x):
            return module(carry, x)
        scan = nn.scan(scan_fn, ....)
        carry = get_carry()
        carry, output = scan(module, carry, x)
        return output

You can simplify it to:

class Model(nn.Module)
    @nn.compact
    def __call__(self, x):
        module = submodule()
        scan = nn.scan(type(module).__call__, ....)
        carry = get_carry()
        carry, output = scan(module, carry, x)
        return output

[cgarciae]

This assumes that module has this signature:

def __call__(self, carry, inputs) -> Tuple[Carry, Output]

[YouJiacheng]

Thanks for your reply! Sorry for my XY problem: I actually want to pass submodule to a library function which uses jax.lax.scan internally.