Add int4/uint4 support in jaxtyping. (#174)

* Add int4/uint4 support in jaxtyping.

* Fix a typo and update api docs.

docs/api/array.md

@@ -66,9 +66,9 @@ The dtype should be any one of (all imported from `jaxtyping`):
                 - Of particular precision: `Complex64`, `Complex128`
         - Any integer or unsigned intger: `Integer`
             - Any unsigned integer: `UInt`
-                - Of particular precision: `UInt8`, `UInt16`, `UInt32`, `UInt64`
+                - Of particular precision: `UInt4`, `UInt8`, `UInt16`, `UInt32`, `UInt64`
             - Any signed integer: `Int`
-                - Of particular precision: `Int8`, `Int16`, `Int32`, `Int64`
+                - Of particular precision: `Int4`, `Int8`, `Int16`, `Int32`, `Int64`
         - Any floating, integer, or unsigned integer: `Real`.
 
 Unless you really want to force a particular precision, then for most applications you should probably allow any floating-point, any integer, etc. That is, use

jaxtyping/__init__.py

@@ -85,6 +85,7 @@ if typing.TYPE_CHECKING:
         Float64 as Float64,
         Inexact as Inexact,
         Int as Int,
+        Int4 as Int4,
         Int8 as Int8,
         Int16 as Int16,
         Int32 as Int32,
@@ -95,6 +96,7 @@ if typing.TYPE_CHECKING:
         Real as Real,
         Shaped as Shaped,
         UInt as UInt,
+        Uint4 as Uint4,
         UInt8 as UInt8,
         UInt16 as UInt16,
         UInt32 as UInt32,
@@ -113,6 +115,7 @@ else:
         Float64 as Float64,
         Inexact as Inexact,
         Int as Int,
+        Int4 as Int4,
         Int8 as Int8,
         Int16 as Int16,
         Int32 as Int32,
@@ -122,6 +125,7 @@ else:
         Real as Real,
         Shaped as Shaped,
         UInt as UInt,
+        UInt4 as UInt4,
         UInt8 as UInt8,
         UInt16 as UInt16,
         UInt32 as UInt32,

jaxtyping/_array_types.py

@@ -652,10 +652,12 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
 _prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
+_uint4 = "uint4"
 _uint8 = "uint8"
 _uint16 = "uint16"
 _uint32 = "uint32"
 _uint64 = "uint64"
+_int4 = "int4"
 _int8 = "int8"
 _int16 = "int16"
 _int32 = "int32"
@@ -681,10 +683,12 @@ def _make_dtype(_dtypes, name):
     return _Cls
 
 
+UInt4 = _make_dtype(_uint4, "UInt4")
 UInt8 = _make_dtype(_uint8, "UInt8")
 UInt16 = _make_dtype(_uint16, "UInt16")
 UInt32 = _make_dtype(_uint32, "UInt32")
 UInt64 = _make_dtype(_uint64, "UInt64")
+Int4 = _make_dtype(_int4, "Int4")
 Int8 = _make_dtype(_int8, "Int8")
 Int16 = _make_dtype(_int16, "Int16")
 Int32 = _make_dtype(_int32, "Int32")
@@ -697,8 +701,8 @@ Complex64 = _make_dtype(_complex64, "Complex64")
 Complex128 = _make_dtype(_complex128, "Complex128")
 
 bools = [_bool, _bool_]
-uints = [_uint8, _uint16, _uint32, _uint64]
-ints = [_int8, _int16, _int32, _int64]
+uints = [_uint4, _uint8, _uint16, _uint32, _uint64]
+ints = [_int4, _int8, _int16, _int32, _int64]
 floats = [_bfloat16, _float16, _float32, _float64]
 complexes = [_complex64, _complex128]
 

jaxtyping/_indirection.py

@@ -32,6 +32,7 @@ from typing import (
     Annotated as Float64,  # noqa: F401
     Annotated as Inexact,  # noqa: F401
     Annotated as Int,  # noqa: F401
+    Annotated as Int4,  # noqa: F401
     Annotated as Int8,  # noqa: F401
     Annotated as Int16,  # noqa: F401
     Annotated as Int32,  # noqa: F401
@@ -42,6 +43,7 @@ from typing import (
     Annotated as Real,  # noqa: F401
     Annotated as Shaped,  # noqa: F401
     Annotated as UInt,  # noqa: F401
+    Annotated as UInt4,  # noqa: F401
     Annotated as UInt8,  # noqa: F401
     Annotated as UInt16,  # noqa: F401
     Annotated as UInt32,  # noqa: F401

test/test_array.py

@@ -67,6 +67,7 @@ def test_dtypes():
         Float64,
         Inexact,
         Int,
+        Int4,
         Int8,
         Int16,
         Int32,
@@ -74,6 +75,7 @@ def test_dtypes():
         Num,
         Shaped,
         UInt,
+        UInt4,
         UInt8,
         UInt16,
         UInt32,
@@ -125,7 +127,9 @@ def test_any_dtype(jaxtyp, typecheck, getkey):
 
     g(jr.normal(getkey(), (3, 4)))
     g(jnp.array([[True, False]]))
+    g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int4))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int8))
+    g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint4))
     g(jnp.array([[1, 2], [3, 4]], dtype=jnp.uint16))
     g(jr.normal(getkey(), (3, 4), dtype=jnp.complex128))
     g(jr.normal(getkey(), (3, 4), dtype=jnp.bfloat16))