[XLA:Python] Adding xla::PrimitiveType <-> numpy.dtype conversions to the library for internal debugging tools.

PiperOrigin-RevId: 620068167

Author: wrengr

third_party/xla/xla/python/tools/BUILD

@@ -63,6 +63,7 @@ tsl_pybind_extension(
         "//xla:literal",
         "//xla:xla_data_proto_cc",
         "//xla/python:logging",
+        "//xla/python:nb_numpy",
         "//xla/python:types",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

third_party/xla/xla/python/tools/_types.cc

@@ -25,6 +25,7 @@ limitations under the License.
 #include "pybind11_protobuf/native_proto_caster.h"  // from @pybind11_protobuf
 #include "xla/literal.h"
 #include "xla/python/logging.h"
+#include "xla/python/nb_numpy.h"
 #include "xla/python/types.h"
 #include "xla/xla_data.pb.h"
 // NOTE: The tsl-numpy header forbids importing the actual NumPy arrayobject.h
@@ -59,6 +60,19 @@ absl::StatusOr<py::object> MakeNdarray(const xla::LiteralProto& proto) {
   // Convert `nb::object` into `py::object`.
   return py::reinterpret_steal<py::object>(nbobj.release().ptr());
 }
+
+// Partial reversion of cl/617156835, until we can get the proto-casters
+// (and hence the extension) switched over to nanobind.
+// TODO(wrengr): Or can we mix `{py,nb}::module_::def` calls??
+absl::StatusOr<xla::PrimitiveType> DtypeToEtype(const py::dtype& py_d) {
+  auto nb_d = nb::borrow<xla::nb_dtype>(py_d.ptr());
+  return xla::DtypeToPrimitiveType(nb_d);
+}
+
+absl::StatusOr<py::dtype> EtypeToDtype(xla::PrimitiveType p) {
+  TF_ASSIGN_OR_RETURN(xla::nb_dtype nb_d, xla::PrimitiveTypeToNbDtype(p));
+  return py::reinterpret_steal<py::dtype>(nb_d.release().ptr());
+}
 }  // namespace
 
 // NOTE: It seems insurmountable to get "native_proto_caster.h" to work
@@ -98,7 +112,8 @@ PYBIND11_MODULE(_types, py_m) {
   py::module_::import("ml_dtypes");
 
   // Ensure that tsl-numpy initializes datastructures of the actual-NumPy
-  // implementation, and does whatever else tsl-numpy needs.
+  // implementation, and does whatever else tsl-numpy needs.  This is
+  // also necessary for using the `xla::nb_dtype` type.
   tsl::ImportNumpy();
 
   // Declare that C++ can `nb::cast` from `std::shared_ptr<xla::Literal>`
@@ -124,5 +139,21 @@ PYBIND11_MODULE(_types, py_m) {
     of tuples with leaves being `numpy.ndarray` views of array-shaped
     sub-literals.
   )pbdoc");
+
+  // This method name is based on `xla_client.dtype_to_etype`.
+  // NOTE: `xla_client` uses a Python class wrapping the protobuf-enum,
+  // rather than using the protobuf-enum directly.  See the module docstring
+  // in "types.py" for more explanation on why.
+  py_m.def("dtype_to_etype", &DtypeToEtype, py::arg("dtype").none(false),
+           py::pos_only(), R"pbdoc(
+    Converts `numpy.dtype` into
+    `tensorflow.compiler.xla.xla_data_pb2.PrimitiveType`.
+  )pbdoc");
+
+  py_m.def("etype_to_dtype", &EtypeToDtype, py::arg("ptype").none(false),
+           py::pos_only(), R"pbdoc(
+    Converts `tensorflow.compiler.xla.xla_data_pb2.PrimitiveType` into
+    `numpy.dtype`.
+  )pbdoc");
   // LINT.ThenChange(_types.pyi)
 }

third_party/xla/xla/python/tools/_types.pyi

@@ -20,4 +20,6 @@ from xla import xla_data_pb2
 # LINT.IfChange
 NdarrayTree = Union[np.ndarray, tuple['NdarrayTree', ...]]
 def make_ndarray(proto: xla_data_pb2.LiteralProto, /) -> NdarrayTree: ...
+def dtype_to_etype(dtype: np.dtype, /) -> xla_data_pb2.PrimitiveType: ...
+def etype_to_dtype(ptype: xla_data_pb2.PrimitiveType, /) -> np.dtype: ...
 # LINT.ThenChange(types.py, _types.cc)

third_party/xla/xla/python/tools/types.py

@@ -40,10 +40,12 @@
 
 # NOTE: `import <name> as <name>` is required for names to be exported.
 # See PEP 484 & <https://github.com/google/jax/issues/7570>
-# pylint: disable=g-importing-member,useless-import-alias,unused-import
+# pylint: disable=g-importing-member,useless-import-alias,unused-import,g-multiple-import
 # LINT.IfChange
 from ._types import (
     make_ndarray as make_ndarray,
+    dtype_to_etype as dtype_to_etype,
+    etype_to_dtype as etype_to_dtype,
 )
 # TODO(wrengr): We can't import the `NdarrayTree` defined in the pyi file.
 # So re-defining it here for now.

third_party/xla/xla/python/tools/types_test.py

@@ -148,14 +148,12 @@ class MakeNdarrayValidTest(parameterized.TestCase):
 
   def testHasCorrectDtype(self, proto, arr):
     """Test that the result has the right dtype."""
-    # Silence [unused-argument] warning.
-    del proto
-    # TODO(wrengr): Add pybind for `xla::PrimitiveTypeToDtype`,
-    # so that we can avoid hard-coding the expected np.dtype.
-    # Alternatively, we could use `xla_client.dtype_to_etype` (ideally
-    # after refactoring that into a small library, so we need not pull in
-    # all the rest of xla_client).
-    self.assertEqual(np.float64, arr.dtype)
+    e = proto.shape.element_type
+    d = arr.dtype
+    with self.subTest(msg='etype_to_dtype'):
+      self.assertEqual(types.etype_to_dtype(e), d)
+    with self.subTest(msg='dtype_to_etype'):
+      self.assertEqual(e, types.dtype_to_etype(d))
 
   def testHasCorrectRank(self, proto, arr):
     """Test that the result has the right rank."""