Exposes graphdef for flax models.

[wang2yn84]

Description

Expose the graphdef for nnx models so that upstream have complete information of the model with both graph and state.

Tunix needs to run the model directly without any kernel optimizations to get the accurate log probs. For non flax models, it already returns model that Tunix can run directly.

Tests

Tested with Tunix + vLLM and it's working fine. Also rely on unit tests in this repo to make sure the nothing basic is wrong.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

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

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[hfan]

Is there a CI run to verify no test change is needed?