Add shape-based lazy init to LinenToNNX (prev LinenWrapper)

[IvyZX]

• Renamed wrappers to LinenToNNX and NNXToLinen to minimize confusion.
• Moved state initialization of LinenToNNX to __call__ to realize lazy init. This allows it to be a submodule of an NNX module, which doesn't have input args during initialization.
  • User can use nnx.shaped_init to do a dry run of __call__ and initialize the whole state & full graphdef.
• Made state initialization of LinenToNNX nested & closer to NNX, aka. each VariableState is created for every jax Array, not every collection.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 51 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (31adb00) to head (0139c90).
Report is 140 commits behind head on main.

Files	Patch %	Lines
flax/nnx/nnx/bridge/wrappers.py	0.00%	48 Missing ⚠️
flax/nnx/nnx/bridge/__init__.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff           @@
##            main   #4081    +/-   ##
======================================
  Coverage   0.00%   0.00%            
======================================
  Files        106     108     +2     
  Lines      13582   14045   +463     
======================================
- Misses     13582   14045   +463     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[cgarciae]

Suggested change

	_rngs['params'] = _rngs['default']
	del _rngs['default']
	_rngs['params'] = _rngs.pop('default')

[cgarciae]

Now that I think about it, we have to make Rngs implement MutableMapping for either of these to work.

[IvyZX]

The _rngs here is a dict instead of Rngs class, so this already works.

[cgarciae]

In this function we should leverage the fact that Object._object__state._initializing still exists and set if via something like

def _set_initializing(initializing: bool):
  for _, value in graph.iter_graph(module):
    if isinstance(value, Object):
      value._object__state._initializing = initializing

and use the value of _initializing to choose between init and apply when calling the Linen Modules.

[IvyZX]

Added _set_initializing to the LinenToNNX wrapper.
Note that we can't do check on top level modules' ._object__state._initializing because the top level module might be a pure NNX module with ._object__state._initializing always False.

[cgarciae]

I'd do something like this:

def _set_initializing(module, initializing: bool): 
  for _, value in graph.iter_graph(module): 
    if isinstance(value, Object): 
      value._object__state._initializing = initializing
      
def shaped_init(module: Module, *args, **kwargs):
  """To trigger init of all `LinenToNNX` module variables and return a wholesome state."""
  module = graph.clone(module) # create a copy
  _set_initializing(module, True)
  assert callable(module)
  try:
    _ = module(*args, **kwargs)
  finally:
    _set_initializing(module, False)
  return nnx.split(module)

[IvyZX]

Done, also renamed to lazy_init as discussed offline.

[cgarciae]

Use self._object__state.initializing instead, see above.

Suggested change

	if not hasattr(self, 'states'):
	if self._object__state.initializing:

[IvyZX]

Done.

[cgarciae]

We need to generate new keys so Linen Modules get new RNG state every time.

Suggested change

	{name: stream.key.raw_value for name, stream in rngs.items()}
	{name: stream() for name, stream in rngs.items()}

[IvyZX]

Done.

[cgarciae]

we could use init_with_output to avoid calling forward twice

[IvyZX]

Done.