NNXWrapper
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Hi,
I have a large library that we decided to build on top of `flax.linen` several years ago.
I'd like now to begin testing `nnx`. However, given the size of the repo and people using it, I cannot change everything at once over to nnx, instead I would like to keep using `linen-style` code for a while, and allowing users to use models defined with nnx inside of our library.
In brief, the way we use modules right now is
```python
model = LinenModel(...)
model_state, parameters = fcore.pop(model.init(jax.random.key(1), ...), "params")
...
# jit boundary
variables ={"params": parameters, **model_state}
model.apply(variables, inoputs...)
```
I tried to use `nnx.split` to this end, but the way it works, returning a special object and not a simple dictionary, makes it impossible to have this approach work fine.
By inspecting `nnx.compat/bridge` I see that you have several utilities to use linen layers within nnx, but it is unclear to me how to do the opposite.
It seems that `nnx.bridge.NNXWrapper` should do that, but it is unfinished, while it is not clear to me how to use `nnx.Module`..
Is there anything I can use?
Contributor guide
Assessment
This issue has not been assessed yet.