Update JAX API usage to latest version

[ReNothingg]

We need to update the JAX API usage across the codebase to use the latest stable versions.

Changes Required

1. Replace jax.tree_util with jax.tree:

• Update all imports
• Replace all usages of tree_util functions

2. Update pytree registration:

• Use register_pytree_with_keys instead of register_pytree_node
• Update flattening/unflattening functions to use new key types
• Update registration calls

3. Update tree traversal operations:

• Replace tree_map with new API version
• Update tree manipulation functions
• Modernize tree-related utilities

Files to Modify

Key files that need updates:

• axlearn/common/struct.py
• axlearn/common/utils.py
• axlearn/common/metrics.py
• axlearn/common/learner.py
• Other files using JAX tree operations

Implementation Details

1. For each file:

• Scan for JAX tree API usage
• Update imports
• Replace deprecated functions
• Update function signatures
• Add type hints where missing

2. Testing:

• Run all tests with latest JAX
• Verify no regressions
• Check backward compatibility
• Add new tests if needed

Success Criteria

• All tests pass with latest JAX version
• No functionality changes
• Clean deprecation warnings
• Improved type safety
• Backward compatible changes

[ReNothingg]

Hi, the workflows need approval to run (GitHub Actions are pending). Can someone with write access approve and run them? @ruomingp pls

[apghml]

Some of these changes have the potential to break things and don't seem to be necessary, as @changlan mentioned. Could you explain for every change, why it is necessary? Also, please do not mark comments as resolved yourself. To streamline reviewing, we only have PR reviewers mark comments as resolved.

[apghml]

Also please resolve any merge conflicts.

[github-actions]

This pull request has been automatically marked as stale because it has been inactive for 60 days. It will be closed in 7 days if no further activity occurs. If you would like to continue working on this, please remove the stale label or leave a comment.