In progress experimention. Add StringDType to JAX's supported types.

PiperOrigin-RevId: 707662268

jax/_src/dtypes.py

@@ -478,9 +478,27 @@ def _issubdtype_cached(a: type | np.dtype | ExtendedDType,
     np.dtype('complex64'),
     np.dtype('complex128'),
 ]
-_jax_types = _bool_types + _int_types + _float_types + _complex_types
-_jax_dtype_set = {float0, *_bool_types, *_int_types, *_float_types, *_complex_types}
 
+_string_types: list[JAXType] = []
+try:
+  import numpy.dtypes as np_dtypes
+
+  if hasattr(np_dtypes, 'StringDType'):
+    _string_types: list[JAXType] = [np_dtypes.StringDType()]  # type: ignore
+except ImportError:
+  np_dtypes = None  # type: ignore
+
+_jax_types = (
+    _bool_types + _int_types + _float_types + _complex_types + _string_types
+)
+_jax_dtype_set = {
+    float0,
+    *_bool_types,
+    *_int_types,
+    *_float_types,
+    *_complex_types,
+    *_string_types,
+}
 
 _dtype_kinds: dict[str, set] = {
   'bool': {*_bool_types},

tests/dtypes_test.py

@@ -26,6 +26,12 @@
 
 import numpy as np
 
+try:
+  import numpy.dtypes as np_dtypes
+except ImportError:
+  np_dtypes = None  # type: ignore
+
+
 import jax
 from jax import numpy as jnp
 from jax._src import earray
@@ -770,16 +776,21 @@ def g(x):
 
 
 class TestPromotionTables(jtu.JaxTestCase):
+  # Not all types are promotable. For example, currently StringDType is not
+  # promotable.
+  promotable_types = [
+      x for x in dtypes._jax_types if not isinstance(x, np_dtypes.StringDType)
+  ]
 
   @parameterized.named_parameters(
       {"testcase_name": f"_{jaxtype=}", "jaxtype": jaxtype}
-      for jaxtype in dtypes._jax_types + dtypes._weak_types)
+      for jaxtype in promotable_types + dtypes._weak_types)
   def testJaxTypeFromType(self, jaxtype):
     self.assertIs(dtypes._jax_type(*dtypes._dtype_and_weaktype(jaxtype)), jaxtype)
 
   @parameterized.named_parameters(
       {"testcase_name": f"_{jaxtype=}", "jaxtype": jaxtype}
-      for jaxtype in dtypes._jax_types + dtypes._weak_types)
+      for jaxtype in promotable_types + dtypes._weak_types)
   def testJaxTypeFromVal(self, jaxtype):
     try:
       val = jaxtype(0)
@@ -789,7 +800,7 @@ def testJaxTypeFromVal(self, jaxtype):
 
   @parameterized.named_parameters(
       {"testcase_name": f"_{dtype=}", "dtype": dtype}
-      for dtype in dtypes._jax_types)
+      for dtype in promotable_types)
   def testJaxTypeWeak(self, dtype):
     jax_type = dtypes._jax_type(dtype, weak_type=True)
     if dtypes.issubdtype(jax_type, np.complexfloating):

tests/export_test.py

@@ -48,6 +48,11 @@
 
 import numpy as np
 
+try:
+  import numpy.dtypes as np_dtypes
+except ImportError:
+  np_dtypes = None  # type: ignore
+
 # ruff: noqa: F401
 try:
   import flatbuffers
@@ -414,7 +419,6 @@ def f(x1, x2):
     self.assertEqual(tree_util.tree_structure(res2),
                      tree_util.tree_structure(res))
 
-
   def test_error_wrong_intree(self):
     def f(a_b_pair, *, c):
       return jnp.sin(a_b_pair[0]) + jnp.cos(a_b_pair[1]) + c
@@ -1009,6 +1013,12 @@ def f_jax(x):  # x: bool[b]
       for dtype in dtypes._jax_types if dtype != np.dtype("bool")
   ])
   def test_poly_numeric_dtypes(self, dtype=np.int32):
+    if hasattr(np_dtypes, "StringDType") and isinstance(
+        dtype, np_dtypes.StringDType
+    ):
+      self.skipTest(
+          "StringDType is not a numeric type"
+      )  # TODO(jmudigonda): revisit.
     if str(dtype) in {"float8_e4m3b11fnuz",
                       "float8_e4m3fnuz",
                       "float8_e5m2fnuz",
@@ -1624,7 +1634,6 @@ def test_multi_platform_unknown_platform(self):
                        platforms=("tpu", "cpu", "cuda", "other"))(x)
     self.assertEqual(exp.platforms, ("tpu", "cpu", "cuda", "other"))
 
-
   def test_multi_platform_with_donation(self):
     f = jax.jit(jnp.sin, donate_argnums=(0,))
     x = np.arange(3, dtype=np.float32)