Experiment with making float/int types inherit from np.floating/np.in…

…teger.

PiperOrigin-RevId: 719426592

ml_dtypes/_src/custom_float.h

@@ -850,7 +850,7 @@ bool RegisterFloatDtype(PyObject* numpy) {
   // bases must be a tuple for Python 3.9 and earlier. Change to just pass
   // the base type directly when dropping Python 3.9 support.
   Safe_PyObjectPtr bases(
-      PyTuple_Pack(1, reinterpret_cast<PyObject*>(&PyGenericArrType_Type)));
+      PyTuple_Pack(1, reinterpret_cast<PyObject*>(&PyFloatingArrType_Type)));
   PyObject* type =
       PyType_FromSpecWithBases(&CustomFloatType<T>::type_spec, bases.get());
   if (!type) {

ml_dtypes/_src/dtypes.cc

@@ -398,10 +398,10 @@ bool Initialize() {
     return false;
   }
 
-  if (!RegisterIntNDtype<int2>(numpy.get()) ||
-      !RegisterIntNDtype<uint2>(numpy.get()) ||
-      !RegisterIntNDtype<int4>(numpy.get()) ||
-      !RegisterIntNDtype<uint4>(numpy.get())) {
+  if (!RegisterIntNDtype<int2>(numpy.get(), /* is_signed= */ true) ||
+      !RegisterIntNDtype<uint2>(numpy.get(), /* is_signed= */ false) ||
+      !RegisterIntNDtype<int4>(numpy.get(), /* is_signed= */ true) ||
+      !RegisterIntNDtype<uint4>(numpy.get(), /* is_signed= */ false)) {
     return false;
   }
 

ml_dtypes/_src/intn_numpy.h

@@ -765,11 +765,13 @@ bool RegisterIntNUFuncs(PyObject* numpy) {
 }
 
 template <typename T>
-bool RegisterIntNDtype(PyObject* numpy) {
+bool RegisterIntNDtype(PyObject* numpy, bool is_signed) {
   // bases must be a tuple for Python 3.9 and earlier. Change to just pass
   // the base type directly when dropping Python 3.9 support.
-  Safe_PyObjectPtr bases(
-      PyTuple_Pack(1, reinterpret_cast<PyObject*>(&PyGenericArrType_Type)));
+  Safe_PyObjectPtr bases(PyTuple_Pack(
+      1, is_signed
+             ? (reinterpret_cast<PyObject*>(&PySignedIntegerArrType_Type))
+             : (reinterpret_cast<PyObject*>(&PyUnsignedIntegerArrType_Type))));
   PyObject* type =
       PyType_FromSpecWithBases(&IntNTypeDescriptor<T>::type_spec, bases.get());
   if (!type) {

ml_dtypes/tests/custom_float_test.py

@@ -590,6 +590,10 @@ def testArgminOnPositiveInfinity(self, float_type):
   def testDtypeFromString(self, float_type):
     assert np.dtype(float_type.__name__) == np.dtype(float_type)
 
+  def testIssubdtype(self, float_type):
+    self.assertTrue(np.issubdtype(float_type, np.floating))
+    self.assertTrue(np.issubdtype(np.dtype(float_type), np.floating))
+
 
 BinaryOp = collections.namedtuple("BinaryOp", ["op"])
 
@@ -713,7 +717,7 @@ def testDeepCopyDoesNotAlterHash(self, float_type):
   def testArray(self, float_type):
     x = np.array([[1, 2, 4]], dtype=float_type)
     self.assertEqual(float_type, x.dtype)
-    self.assertEqual("[[1 2 4]]", str(x))
+    self.assertEqual("[[1. 2. 4.]]", str(x))
     np.testing.assert_equal(x, x)
     numpy_assert_allclose(x, x, float_type=float_type)
     self.assertTrue((x == x).all())

ml_dtypes/tests/intn_test.py

@@ -245,6 +245,11 @@ def testCanCast(self, a, b):
         ((a, b) in allowed_casts), np.can_cast(a, b, casting="safe")
     )
 
+  @parameterized.product(scalar_type=INTN_TYPES)
+  def testIssubdtype(self, scalar_type):
+    self.assertTrue(np.issubdtype(scalar_type, np.integer))
+    self.assertTrue(np.issubdtype(np.dtype(scalar_type), np.integer))
+
 
 # Tests for the Python scalar type
 @multi_threaded(num_workers=3, skip_tests=["testBinaryUfuncs"])
@@ -274,7 +279,7 @@ def testDeepCopyDoesNotAlterHash(self, scalar_type):
   def testArray(self, scalar_type):
     if scalar_type == int2:
       x = np.array([[-2, 1, 0, 1]], dtype=scalar_type)
-      self.assertEqual("[[-2 1 0 1]]", str(x))
+      self.assertEqual("[[-2  1  0  1]]", str(x))
     else:
       x = np.array([[1, 2, 3]], dtype=scalar_type)
       self.assertEqual("[[1 2 3]]", str(x))