google / google/flax

NNXWrapper

Open
#4,088 2 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.