Swapping selected layers with a different layer.

Hello!

I am pretty new to Jax/Flax. I was trying to implement to LoRA, although I could modify the state dict, would it be possible to modify the layers to accommodate the additional matrices on run time? Or, is there any other alternate approach to this, so that I could freeze the original weights and just keep the $A$ and $B$ matrices trainable?

I am able to modify the state dict by doing something like this:
```python
import jax
import jax.numpy as jnp

PRNGKey = jnp.ndarray

def get_param(path, param, rank: int, rng: PRNGKey):
if len(param.shape) == 1:
return param

a_dim, b_dim = param.shape
return {
"a": jax.random.normal(rng, shape=(a_dim, rank)),
"b": jnp.zeros(shape=(rank, b_dim)),
"w": param
}

new_params = jax.tree_util.tree_map_with_path(get_param, params, is_leaf=None)
```

Is there any maneuver I could do? Or, would I have to make drastic moves like re-defining the architecture?

Any suggestions would help!

Thank you!