Transparent Sub-Module Scope

[joeryjoery]

Hi, I want to write a Module B that depends on an arbitrary sub-module tree A, but I don't want to the top-level of the sub-modules A to be included in the "params" container. Is this possible in Flax?

In other words, I would like to have control over the parameter scope to prevent unnecesary nesting. Unfortunately, I couldn't understand how to do this with pack/ lift based on the docs.

I want to use this for writing module compositions, e.g., for writing loss functions for neural networks, without having to worry about variable or rng management.

To illustrate what I want to achieve:

import flax.linen as nn

import jax
import jax.numpy as jnp


class A(nn.Module):
    leaf: nn.Module = nn.Dense(10)
    
    @nn.compact
    def __call__(self, x):
        return self.leaf(x)
    

class B(nn.Module):
    a: A
    
    @nn.compact    
    def __call__(self, x):
        return self.a.leaf(x)
    

myA = A()
myB = B(myA)

vA = myA.init(jax.random.PRNGKey(0), jnp.zeros(3))
jax.tree_map(jnp.shape, vA)
>> {'params': {'leaf': {'bias': (10,), 'kernel': (3, 10)}}}  # Desired

vB = myB.init(jax.random.PRNGKey(0), jnp.zeros(3))
jax.tree_map(jnp.shape, vB)
>> {'params': {'a': {'leaf': {'bias': (10,), 'kernel': (3, 10)}}}}  # Actual

The reason why I desire the second param-container, without the 'a' scope, is that I want to initialize from A and be able to use the same parameter container in B:

# This throws an error because 'a' is not in the vA container.
myB.apply(vA, jnp.zeros(3))  
>> ScopeParamNotFoundError                   Traceback (most recent call last)
Cell In[18], line 4
      1 v = myA.init(jax.random.PRNGKey(0), jnp.zeros(3))
      2 jax.tree_map(jnp.shape, v)
----> 4 myB.apply(v, jnp.zeros(3))

Workaround Pattern

I was able to get around my problem by passing custom methods inside the apply function of myA:

import flax.linen as nn

import jax
import jax.numpy as jnp


class A(nn.Module):
    leaf: nn.Module = nn.Dense(10)
    
    @nn.compact
    def __call__(self, x):
        return self.leaf(x)

    
def dummy(module: A, x):
    return module.leaf(x) + 1

    
myA = A()
vA = myA.init(jax.random.PRNGKey(0), jnp.zeros(3))

myA.apply(vA, jnp.zeros(3))
>> Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32)

myA.apply(vA, jnp.zeros(3), method=dummy)
>> Array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)

Or with dataclasses,

@flax.struct.dataclass
class B:
    a: A
    
    def __call__(self, v, x):
        return self.a.apply(v, x, method=self.dummy)
    
    def dummy(self, module, x):
        return module.leaf(x) + 1

I'm not sure if these patterns are ideal, the former uses a method dummy that does not belong to A and the latter breaks consistency with the Module interface.

Ideally, I could implement a Loss Module dependent on A where the call myLoss.apply(vars_from_A, ...) just works, does anyone know how to achieve this?