Merge pull request #20816 from superbobry:int4

PiperOrigin-RevId: 626131730

Author: jax authors

jax/_src/test_util.py

@@ -1289,15 +1289,17 @@ def integer(self):
 
   @_cached_property
   def all_integer(self):
-    return self.supported([np.int8, np.int16, np.int32, np.int64])
+    return self.supported([
+        _dtypes.int4, np.int8, np.int16, np.int32, np.int64])
 
   @_cached_property
   def unsigned(self):
     return self.supported([np.uint32, np.uint64])
 
   @_cached_property
   def all_unsigned(self):
-    return self.supported([np.uint8, np.uint16, np.uint32, np.uint64])
+    return self.supported([
+        _dtypes.uint4, np.uint8, np.uint16, np.uint32, np.uint64])
 
   @_cached_property
   def complex(self):

jax/tools/jax_to_ir.py

@@ -238,12 +238,14 @@ def parse_shape_str(s):
     shape = ()
   return jax.core.ShapedArray(shape, dtype)
 
-_DT = {'pred': jnp.bool_,
-       'u8': jnp.uint8, 'u16': jnp.uint16, 'u32': jnp.uint32, 'u64': jnp.uint64,
-       's8': jnp.int8, 's16': jnp.int16, 's32': jnp.int32, 's64': jnp.int64,
-       'bf16': jnp.bfloat16,
-       'f16': jnp.float16, 'f32': jnp.float32, 'f64': jnp.float64,
-       'c64': jnp.complex64, 'c128': jnp.complex128}
+_DT = {
+    'pred': jnp.bool_,
+    'u4': jnp.uint4, 'u8': jnp.uint8, 'u16': jnp.uint16, 'u32': jnp.uint32, 'u64': jnp.uint64,
+    's4': jnp.int4, 's8': jnp.int8, 's16': jnp.int16, 's32': jnp.int32, 's64': jnp.int64,
+    'bf16': jnp.bfloat16,
+    'f16': jnp.float16, 'f32': jnp.float32, 'f64': jnp.float64,
+    'c64': jnp.complex64, 'c128': jnp.complex128
+}
 _SHAPE_RE = re.compile(f"^({'|'.join(_DT)})\\[\\s*(\\d*[\\s*,\\d+]*)\\s*\\]$")
 
 

tests/dtypes_test.py

@@ -67,8 +67,8 @@
 all_dtypes = (bool_dtypes + signed_dtypes + unsigned_dtypes + float_dtypes +
               complex_dtypes)
 
-scalar_types = [jnp.bool_, jnp.int8, jnp.int16, jnp.int32, jnp.int64,
-                jnp.uint8, jnp.uint16, jnp.uint32, jnp.uint64,
+scalar_types = [jnp.bool_, jnp.int4, jnp.int8, jnp.int16, jnp.int32, jnp.int64,
+                jnp.uint4, jnp.uint8, jnp.uint16, jnp.uint32, jnp.uint64,
                 jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64,
                 jnp.complex64, jnp.complex128]
 
@@ -94,6 +94,7 @@
 _EXPECTED_CANONICALIZE_X32[np.longlong] = np.int32
 
 UINT_DTYPES = {
+  4: jnp.uint4,
   8: np.uint8,
   16: np.uint16,
   32: np.uint32,
@@ -284,13 +285,15 @@ def testIsSubdtype(self):
       self.assertTrue(dtypes.issubdtype(np.dtype(t).type, t))
       self.assertTrue(dtypes.issubdtype(t, np.dtype(t).type))
       self.assertTrue(dtypes.issubdtype(t, np.dtype(t)))
-      if t != jnp.bfloat16:
-        for category in [np.generic, jnp.inexact, jnp.integer, jnp.signedinteger,
-                         jnp.unsignedinteger, jnp.floating, jnp.complexfloating]:
-          self.assertEqual(dtypes.issubdtype(t, category),
-                           np.issubdtype(np.dtype(t).type, category))
-          self.assertEqual(dtypes.issubdtype(t, category),
-                           np.issubdtype(np.dtype(t).type, category))
+      if t in [jnp.int4, jnp.uint4, jnp.bfloat16]:
+        # These dtype have no equivalent in NumPy.
+        continue
+      for category in [np.generic, jnp.inexact, jnp.integer, jnp.signedinteger,
+                       jnp.unsignedinteger, jnp.floating, jnp.complexfloating]:
+        self.assertEqual(dtypes.issubdtype(t, category),
+                         np.issubdtype(np.dtype(t).type, category))
+        self.assertEqual(dtypes.issubdtype(t, category),
+                         np.issubdtype(np.dtype(t).type, category))
 
   def testIsSubdtypeExtended(self):
     self.assertTrue(dtypes.issubdtype(dtypes.extended, dtypes.extended))

tests/jax_to_ir_test.py

@@ -119,10 +119,12 @@ def test_parse_shape_str(self):
     self.assertParsedShape('f32[]', [], jnp.float32)
     self.assertParsedShape('f32[1,2,3]', [1, 2, 3], jnp.float32)
     self.assertParsedShape('pred[1]', [1], jnp.bool_)
+    self.assertParsedShape('s4[1]', [1], jnp.int4)
     self.assertParsedShape('s8[1]', [1], jnp.int8)
     self.assertParsedShape('s16[1]', [1], jnp.int16)
     self.assertParsedShape('s32[1]', [1], jnp.int32)
     self.assertParsedShape('s64[1]', [1], jnp.int64)
+    self.assertParsedShape('u4[1]', [1], jnp.uint4)
     self.assertParsedShape('u8[1]', [1], jnp.uint8)
     self.assertParsedShape('u16[1]', [1], jnp.uint16)
     self.assertParsedShape('u32[1]', [1], jnp.uint32)

tests/lax_numpy_test.py

@@ -195,11 +195,13 @@ def testUnstack(self, shape, axis, dtype):
 
 
   @parameterized.parameters(
-    [dtype for dtype in [jnp.bool, jnp.uint8, jnp.uint16, jnp.uint32,
-                         jnp.uint64, jnp.int8, jnp.int16, jnp.int32, jnp.int64,
-                         jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64,
-                         jnp.complex64, jnp.complex128]
-     if dtype == dtypes.canonicalize_dtype(dtype)])
+      [dtype for dtype in [
+          jnp.bool,
+          jnp.uint4, jnp.uint8, jnp.uint16, jnp.uint32, jnp.uint64,
+          jnp.int4, jnp.int8, jnp.int16, jnp.int32, jnp.int64,
+          jnp.bfloat16, jnp.float16, jnp.float32, jnp.float64,
+          jnp.complex64, jnp.complex128]
+       if dtype == dtypes.canonicalize_dtype(dtype)])
   def testDtypeWrappers(self, dtype):
     arr = dtype(0)
     self.assertIsInstance(arr, jax.Array)