Skip to content

Support bool and float16 dtypes and validate sharding in persistence helper. #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Support bool and float16 dtypes and validate sharding in persistence helper. #370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions pathwaysutils/persistence/helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions pathwaysutils/test/persistence_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Tests for the persistence module."""

import datetime
from unittest import mock

from absl.testing import absltest
import jax
Expand Down Expand Up @@ -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()
Loading