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

PiperOrigin-RevId: 707662268
jax-ml · Dec 20, 2024 · f9d2b96 · f9d2b96
1 parent 4216f8f
commit f9d2b96
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
@@ -476,9 +476,15 @@ 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}
-
+try:
+  import numpy.dtypes as np_dtypes
+  _string_types: list[JAXType] = [np_dtypes.StringDType()]
+except ImportError:
+  _string_types: list[JAXType] = []
+  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},