Add memory_kind to the PositionalShardingMetadata in `test_convert_…

…between_jax_positional_sharding_and_sharding_metadata` test PiperOrigin-RevId: 660548319
google · Aug 7, 2024 · a301fc8 · a301fc8
1 parent 65416bb
commit a301fc8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/checkpoint/orbax/checkpoint/metadata/sharding_test.py b/checkpoint/orbax/checkpoint/metadata/sharding_test.py
@@ -55,7 +55,9 @@ def test_convert_between_jax_positional_sharding_and_sharding_metadata(
         [1, -1]
     )
     expected_positional_sharding_metadata = (
-        sharding_metadata.PositionalShardingMetadata(jax_sharding.shape)
+        sharding_metadata.PositionalShardingMetadata(
+            jax_sharding.shape, jax_sharding.memory_kind
+        )
     )
     converted_positional_sharding_metadata = (
         sharding_metadata.from_jax_sharding(jax_sharding)