Skip to content

Comments

NNX TrainState and unit tests#3024

Open
xibinliu wants to merge 1 commit intoxibin/nnx_1from
xibin/nnx_2
Open

NNX TrainState and unit tests#3024
xibinliu wants to merge 1 commit intoxibin/nnx_1from
xibin/nnx_2

Conversation

@xibinliu
Copy link
Collaborator

@xibinliu xibinliu commented Jan 28, 2026

Description

Start with a short description of what the PR does and how this is a change from
the past.

A TrainState for NNX, which includes model and optimizer Unit tests include checkpoint tests:

  • restore a saved state
  • convert linen TrainState to NNX TrainState
  • Parameter only restore (no opt_state)

Tests

All unit tests passed.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@xibinliu
NNX TrainState and unit tests
d71b276 
A TrainState for NNX, which includes model and optimizer
Unit tests include checkpoint tests:
- restore a saved state
- convert linen TrainState to NNX TrainState
- Parameter only restore (no opt_state)
@codecov
Copy link

codecov bot commented Jan 28, 2026 •
edited
Loading

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@RissyRan RissyRan marked this pull request as ready for review January 28, 2026 18:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@xibinliu