Support 2D transposed shape alignment in model_creation_utils

[suhithr]

When restoring scanned MoE checkpoints (such as DeepSeek V3), router gate bias (e.g. MoEBiasVar) may be stored on disk with shape (num_experts, num_moe_layers) due to converter scan-axis stacking, whereas Flax NNX defines the parameter as (num_moe_layers, num_experts).

Handle 2D transposed dimensions in _align_checkpoint_to_model_shapes() by transposing the checkpoint array to match target model shapes.

FIXES: b/553744097

Tests

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.

[gemini-code-assist]

Code Review

This pull request adds support for aligning transposed 2D checkpoint arrays to the model shape in _align_checkpoint_to_model_shapes and includes corresponding unit tests. The reviewer suggested simplifying the transposition logic by removing a redundant shape check, as transposing a 2D array with reversed dimensions mathematically guarantees it will match the target model shape.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!