promote nnx to root-level import

[chiamp]

promote nnx to root-level import

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 1 line in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (31adb00) to head (c206b34). Report is 80 commits behind head on main.

Files	Patch %	Lines
flax/__init__.py	0.00%	1 Missing :warning:

Additional details and impacted files

@@          Coverage Diff           @@
##            main   #3961    +/-   ##
======================================
  Coverage   0.00%   0.00%            
======================================
  Files        106     107     +1     
  Lines      13582   13825   +243     
======================================
- Misses     13582   13825   +243     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.