nn.share_scope and initialising with existing parameters

[PhilipVinc]

Is there a simple way to initialise the parameters of a wrapper linen Module, while keeping the parameters of the wrapped module?

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

model = nn.Dense(10)  
model_wrapped = DenseLoRA(model, 2)

variables = model.init(random.key(0), jnp.ones((1, 5))

# get something like
variables_wrapped = model_wrapped.init(random.key(0), jnp.ones((1, 5))
variables_wrapped["params"]["kernel"] = variables["params"]["kernel"]
variables_wrapped["params"]["bias"] = variables["params"]["bias"]

but that works regardless of the wrapped and wrappee structure?