Add unittest for safe cast for (u)int4 type; remove safe cast from in…

…t8 to int4 types PiperOrigin-RevId: 573102837
jax-ml · Oct 13, 2023 · 348fd37 · 348fd37
1 parent baca1aa
commit 348fd37
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 8 deletions.
diff --git a/ml_dtypes/_src/int4_numpy.h b/ml_dtypes/_src/int4_numpy.h
@@ -718,6 +718,10 @@ bool RegisterInt4Casts() {
       return false;
     }
   }
+  if (PyArray_RegisterCanCast(&TypeDescriptor<T>::npy_descr, NPY_HALF,
+                              NPY_NOSCALAR) < 0) {
+    return false;
+  }
   if (PyArray_RegisterCanCast(&TypeDescriptor<T>::npy_descr, NPY_FLOAT,
                               NPY_NOSCALAR) < 0) {
     return false;
@@ -748,14 +752,6 @@ bool RegisterInt4Casts() {
                               TypeDescriptor<T>::Dtype(), NPY_NOSCALAR) < 0) {
     return false;
   }
-  if (PyArray_RegisterCanCast(PyArray_DescrFromType(NPY_UBYTE),
-                              TypeDescriptor<T>::Dtype(), NPY_NOSCALAR) < 0) {
-    return false;
-  }
-  if (PyArray_RegisterCanCast(PyArray_DescrFromType(NPY_BYTE),
-                              TypeDescriptor<T>::Dtype(), NPY_NOSCALAR) < 0) {
-    return false;
-  }
 
   return true;
 }

diff --git a/ml_dtypes/tests/int4_test.py b/ml_dtypes/tests/int4_test.py
@@ -155,6 +155,64 @@ def testBinop(self, scalar_type, op):
           self.assertIsInstance(out, scalar_type)
           self.assertEqual(scalar_type(op(v, w)), out, msg=(v, w))
 
+  @parameterized.product(
+      scalar_type=INT4_TYPES,
+      dtype=[
+          np.float16,
+          np.float32,
+          np.float64,
+          np.int8,
+          np.int16,
+          np.int32,
+          np.int64,
+          np.complex64,
+          np.complex128,
+          np.uint8,
+          np.uint16,
+          np.uint32,
+          np.uint64,
+          np.intc,
+          np.int_,
+          np.longlong,
+          np.uintc,
+          np.ulonglong,
+      ],
+  )
+  def testCanCast(self, scalar_type, dtype):
+    allowed_casts = [
+        (np.bool_, int4),
+        (int4, np.int8),
+        (int4, np.int16),
+        (int4, np.int32),
+        (int4, np.int64),
+        (int4, np.float16),
+        (int4, np.float32),
+        (int4, np.float64),
+        (int4, np.complex64),
+        (int4, np.complex128),
+        (np.bool_, uint4),
+        (uint4, np.int8),
+        (uint4, np.int16),
+        (uint4, np.int32),
+        (uint4, np.int64),
+        (uint4, np.uint8),
+        (uint4, np.uint16),
+        (uint4, np.uint32),
+        (uint4, np.uint64),
+        (uint4, np.float16),
+        (uint4, np.float32),
+        (uint4, np.float64),
+        (uint4, np.complex64),
+        (uint4, np.complex128),
+    ]
+
+    assert ((scalar_type, dtype) in allowed_casts) == np.can_cast(
+        scalar_type, dtype
+    )
+    assert ((dtype, scalar_type) in allowed_casts) == np.can_cast(
+        dtype, scalar_type
+    )
+
 
 # Tests for the Python scalar type
 class ArrayTest(parameterized.TestCase):