[nnx] remove pytreelib

[cgarciae]

What does this PR do?

Removes pytreelib which is no longer used as TrainState is now implemented using flax.struct.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 60.17%. Comparing base (5d9fd9f) to head (4570682).
Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3816      +/-   ##
==========================================
- Coverage   60.86%   60.17%   -0.70%     
==========================================
  Files         105      103       -2     
  Lines       13045    12795     -250     
==========================================
- Hits         7940     7699     -241     
+ Misses       5105     5096       -9     

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