An obsolete module? (flax.core.nn)

[VIVelev]

Hey,

I was wondering, what is exactly the purpose of the module flax.core.nn? I see it is used only in some test cases from tests/core, but not examples or other parts of the flax src code. Is it an out-of-date module that was necessary before, but is now due to be removed? Does it provide any functionality on the user API side?

Thanks!

[avital]

Hi @VIVelev -- this directory contained a nn library implementation directly on top of function core (flax.core), without any module classes. Everything is just a function that takes a scope as a first argument. This was instrumental in our design discussions (and I believe we use these in the flax.core unit tests) but we don't suggest this as a user-facing API.