feat: ported nnx model converter

[mesakhcienet]

Description

This PR introduces scripts to convert Linen model weights into nnx following the changes made to models.py in this pull request (PR #2709).

The changes in PR #2709 result in an incompatibility, causing models trained with the older Linen weights (before PR #2709) to fail loading. These scripts provide a forward-compatibility path.

The added tools are:

1. tools/weight_inspector/convert_linen_checkpoint_to_nnx.py: Converts pre-existing Linen model weights to the current NNX-ported model format.
2. tools/weight_inspector/compare_checkpoint.py: Compares the structure and values between two checkpoint files.

Tests

To ensure the conversion works as expected:

1. Start with a checkpoint directory (/pre-existing-model-path/0/) containing model weights trained before the models.py changes in PR feat: migrate models.py to use nnx #2709.

2. Convert the pre-existing Linen model weights to the new NNX format:

python tools/weight_inspector/convert_linen_checkpoint_to_nnx.py \
  --source_path="/pre-existing-model-path/0/" \
  --output_path="/converted-model/0/"

3. Compare the structure of the newly converted checkpoint with a checkpoint from an NNX model trained after the models.py changes (e.g., /nnx-trained-model/14/items):

python tools/weight_inspector/compare_checkpoint.py \
 --lhs "/converted-model/0/items" \
 --rhs "/nnx-trained-model/14/items"

The comparison script should report no significant structural differences, indicating a successful conversion.
NOTE: For numerical difference, we assume it was caused by different implentation and random generator of Linen and NNX implementation as the main reason.

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.