[linen] add share_scope

[cgarciae]

What does this PR do?

Addsnn.share_scope can be used to share a scope between two Modules. This means that any parameters created by any of the Modules will be added to their common scope. This is useful when you want to wrap a Module and extend its functionality without changing the parameter structure.

import flax.linen as nn
import jax
from jax import numpy as jnp, random

class DenseLoRA(nn.Module):
  base: nn.Dense
  rank: int

  def setup(self):
    nn.share_scope(self, self.base)

  @nn.compact
  def __call__(self, x: jax.Array):
    din, dout = x.shape[-1], self.base.features
    A = self.param('A', nn.zeros_init(), (din, self.rank))
    B = self.param('B', nn.zeros_init(), (self.rank, dout))
    return self.base(x) + x @ A @ B

class Model(nn.Module):
  @nn.compact
  def __call__(self, x: jax.Array):
    base = nn.Dense(10)  # base scope
    return DenseLoRA(base, rank=2)(x)  # reuse base scope

model = Model()

params = model.init(random.key(0), jnp.ones((1, 5)))['params']
list(params['Dense_0'].keys())
# ['A', 'B', 'kernel', 'bias']

[marksandler2]

I would also highlight the difference between

class DenseLora(Dense):
...

vs

class DenseLora(Transparent[Dense]):
...

(The important distinction is that Transparent replaces inheritance with composition, that allows to swap out implementation of inner module with Dense implementation (as long as it has compatible interface), without having to create separate DenseLora for each implementation.

[cgarciae]

@marksandler2 wondering if we should just expose _assign_shared_scope as e.g. nn.share_scope an have users manually do the call instead of introducing the Transparent abstraction? The LoRA example would look like this:

class DenseLoRA(nn.Module):
  inner: nn.Dense
  rank: int

  def setup(self):
    nn.share_scope(self, self.inner)

  @nn.compact
  def __call__(self, x: jax.Array):
    din, dout = x.shape[-1], self.inner.features
    A = self.param('A', nn.zeros_init(), (din, self.rank))
    B = self.param('B', nn.zeros_init(), (self.rank, dout))
    return self.inner(x) + x @ A @ B

class Model(nn.Module):
  @nn.compact
  def __call__(self, x: jax.Array):
    return DenseLoRA(nn.Dense(10), rank=2)(x)

[marksandler2]

Hmm, i like this new proposed function a lot actually. It certainly feels a lot less heavy weight (e.g. we don't need to muck around with Transparent[...] inheritance, and the logic could be made a lot more flexible.

[marksandler2]

Can we merge this?