ensure doctest runs on NNX

[chiamp]

Ensure doctest runs on NNX and fix NNX docstrings.

Fixed some NNX examples and toy examples, but will exclude testing on those until they are all fixed in a future PR.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 61.19%. Comparing base (39ce98f) to head (fbc1fc5).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3797   +/-   ##
=======================================
  Coverage   61.19%   61.19%           
=======================================
  Files         105      105           
  Lines       13148    13148           
=======================================
  Hits         8046     8046           
  Misses       5102     5102           

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

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB