Cache sharding metadata device mesh serialization

[Trighap52]

Summary

Fixes #3494.

This PR deduplicates NamedShardingMetadata device-mesh serialization by caching the DeviceMetadataMesh conversion from a JAX mesh and by caching the serialized dataclasses.asdict output for each metadata mesh.

Details

NamedShardingMetadata.to_serialized_string() previously called dataclasses.asdict(self.device_mesh) for every array metadata entry, which repeatedly deep-copied the same device mesh when many arrays shared one mesh.

This change:

• builds a hashable cache key from the mesh device-id tree instead of relying on jax.sharding.Mesh hashability
• memoizes DeviceMetadataMesh construction for repeated equivalent meshes
• stores the expensive serialized dictionary as a cached_property
• preserves the existing list-shaped device metadata used by serialization and round trips

Validation

• .venv/bin/python checkpoint/orbax/checkpoint/_src/metadata/sharding_test.py
• PATH="$PWD/.venv/bin:$PATH" .venv/bin/pre-commit run --files $(git diff --name-only origin/main)