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()