From 58956f5c77f7f00a606ed687b54e0390cf90f694 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Mon, 26 Oct 2026 00:01:19 -0700 Subject: [PATCH] Support bool and float16 dtypes and validate sharding in persistence helper. PiperOrigin-RevId: 1007052531 --- pathwaysutils/persistence/helper.py | 9 +++++++-- pathwaysutils/test/persistence_test.py | 20 ++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/pathwaysutils/persistence/helper.py b/pathwaysutils/persistence/helper.py index d2bd745..4c5d56e 100644 --- a/pathwaysutils/persistence/helper.py +++ b/pathwaysutils/persistence/helper.py @@ -28,8 +28,12 @@ def dtype_to_xla_primitive_type_str(dtype: np.dtype) -> str: """Converts a numpy dtype to an xla PrimitiveType.""" - if dtype == np.dtype("bfloat16"): + if dtype == np.dtype("bool"): + return "PRED" + elif dtype == np.dtype("bfloat16"): return "BF16" + elif dtype == np.dtype("float16"): + return "F16" elif dtype == np.dtype("float32"): return "F32" elif dtype == np.dtype("float64"): @@ -111,7 +115,8 @@ def get_write_request( ) -> str | Mapping[str, Any]: """Returns a string representation of the plugin program which writes the given jax_array to the given location.""" sharding = jax_array.sharding - assert isinstance(sharding, jax.sharding.Sharding), sharding + if not isinstance(sharding, jax.sharding.Sharding): + raise ValueError(f"Expected jax.sharding.Sharding, got {sharding}") timeout_seconds, timeout_fractional_seconds = divmod( timeout.total_seconds(), 1 diff --git a/pathwaysutils/test/persistence_test.py b/pathwaysutils/test/persistence_test.py index ef1b06f..6ff085b 100644 --- a/pathwaysutils/test/persistence_test.py +++ b/pathwaysutils/test/persistence_test.py @@ -1,6 +1,7 @@ """Tests for the persistence module.""" import datetime +from unittest import mock from absl.testing import absltest import jax @@ -107,6 +108,25 @@ def test_get_bulk_write_request(self): ) self.assertNotEmpty(bulk_write_request) + def test_dtype_to_xla_primitive_type_str_bool_and_float16(self): + self.assertEqual( + helper.dtype_to_xla_primitive_type_str(np.dtype("bool")), "PRED" + ) + self.assertEqual( + helper.dtype_to_xla_primitive_type_str(np.dtype("float16")), "F16" + ) + + def test_get_write_request_invalid_sharding_raises_value_error(self): + fake_array = mock.MagicMock(spec=jax.Array) + fake_array.sharding = "not_a_sharding" + with self.assertRaisesRegex(ValueError, "Expected jax.sharding.Sharding"): + helper.get_write_request( + location_path=self.location, + name=self.name, + jax_array=fake_array, + timeout=self.timeout, + ) + if __name__ == "__main__": absltest.main()