Description
I can't seem to get a NNX module to work with FlaxLayer. Normal flax works fine, but when I create a nnx.Module and called nnx.bridge.ToLinen I get the error at the botttom.
class MyFlax(nn.Module):
@nn.compact
def call(self, x):
return nn.Dense(features=1)(x)
l = MyFlax()
f = keras.layers.FlaxLayer(l)
f(jnp.ones(1))
class MyNnx(nnx.Module):
def init(self, *, rngs: nnx.Rngs):
self.l = nnx.Linear(1, 1, rngs=rngs)
def call(self, **kwargs):
print( kwargs)
if 'inputs' in kwargs:
return self.l(kwargs['inputs'])
MyNnx(rngs=nnx.Rngs(0))(inputs=jnp.ones(1))
l = nnx.bridge.ToLinen(MyNnx(rngs=nnx.Rngs(0)))
f = keras.layers.FlaxLayer(l)
f(inputs=jnp.ones(1))
Cell In[101], line 1
----> 1 f(inputs=jnp.ones(1))
File ~/projects/joe/.venv/lib/python3.11/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback..error_handler(*args, **kwargs)
119 filtered_tb = _process_traceback_frames(e.traceback)
120 # To get the full stack trace, call:
121 # keras.config.disable_traceback_filtering()
--> 122 raise e.with_traceback(filtered_tb) from None
123 finally:
124 del filtered_tb
[... skipping hidden 9 frame]
File ~/projects/joe/.venv/lib/python3.11/site-packages/flax/nnx/bridge/wrappers.py:263, in ToLinen.call(self, *args, **kwargs)
260 # TODO: add lazy_init here in case there's an ToNNX submodule under module.
261 # update linen variables before call module to save initial state
262 self._update_variables(module)
--> 263 out = module(*args, **kwargs)
264 return out
266 # create state
TypeError: 'NoneType' object is not callable