add uint2/int2

following JAX's uint2/int2 data types: https://github.com/jax-ml/jax/blob/main/jax/_src/dtypes.py#L182

Author: mattjj

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: `UInt4`, `UInt8`, `UInt16`, `UInt32`, `UInt64`
+                - Of particular precision: `UInt2`, `UInt4`, `UInt8`, `UInt16`, `UInt32`, `UInt64`
             - Any signed integer: `Int`
-                - Of particular precision: `Int4`, `Int8`, `Int16`, `Int32`, `Int64`
+                - Of particular precision: `Int2`, `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

@@ -66,6 +66,7 @@
         Float64 as Float64,
         Inexact as Inexact,
         Int as Int,
+        Int2 as Int2,
         Int4 as Int4,
         Int8 as Int8,
         Int16 as Int16,
@@ -80,6 +81,7 @@
         ScalarLike as ScalarLike,
         Shaped as Shaped,
         UInt as UInt,
+        UInt2 as UInt2,
         UInt4 as UInt4,
         UInt8 as UInt8,
         UInt16 as UInt16,
@@ -123,6 +125,7 @@
         Float64 as Float64,
         Inexact as Inexact,
         Int as Int,
+        Int2 as Int2,
         Int4 as Int4,
         Int8 as Int8,
         Int16 as Int16,
@@ -134,6 +137,7 @@
         Real as Real,
         Shaped as Shaped,
         UInt as UInt,
+        UInt2 as UInt2,
         UInt4 as UInt4,
         UInt8 as UInt8,
         UInt16 as UInt16,

jaxtyping/_array_types.py

@@ -735,11 +735,13 @@ def __init_subclass__(cls, **kwargs):
 _prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
+_uint2 = "uint2"
 _uint4 = "uint4"
 _uint8 = "uint8"
 _uint16 = "uint16"
 _uint32 = "uint32"
 _uint64 = "uint64"
+_int2 = "int2"
 _int4 = "int4"
 _int8 = "int8"
 _int16 = "int16"
@@ -772,11 +774,13 @@ class _Cls(AbstractDtype):
     return _Cls
 
 
+UInt2 = _make_dtype(_uint2, "UInt2")
 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")
+Int2 = _make_dtype(_int2, "Int2")
 Int4 = _make_dtype(_int4, "Int4")
 Int8 = _make_dtype(_int8, "Int8")
 Int16 = _make_dtype(_int16, "Int16")
@@ -795,8 +799,8 @@ class _Cls(AbstractDtype):
 Complex128 = _make_dtype(_complex128, "Complex128")
 
 bools = [_bool, _bool_]
-uints = [_uint4, _uint8, _uint16, _uint32, _uint64]
-ints = [_int4, _int8, _int16, _int32, _int64]
+uints = [_uint2, _uint4, _uint8, _uint16, _uint32, _uint64]
+ints = [_int2, _int4, _int8, _int16, _int32, _int64]
 float8 = [
     _float8_e4m3b11fnuz,
     _float8_e4m3fn,

jaxtyping/_indirection.py

@@ -37,6 +37,7 @@
     Annotated as Float64,  # noqa: F401
     Annotated as Inexact,  # noqa: F401
     Annotated as Int,  # noqa: F401
+    Annotated as Int2,  # noqa: F401
     Annotated as Int4,  # noqa: F401
     Annotated as Int8,  # noqa: F401
     Annotated as Int16,  # noqa: F401
@@ -48,6 +49,7 @@
     Annotated as Real,  # noqa: F401
     Annotated as Shaped,  # noqa: F401
     Annotated as UInt,  # noqa: F401
+    Annotated as UInt2,  # noqa: F401
     Annotated as UInt4,  # noqa: F401
     Annotated as UInt8,  # noqa: F401
     Annotated as UInt16,  # noqa: F401

test/test_all_importable.py

@@ -18,12 +18,14 @@ def test_all_importable():
         Complex128,  # noqa: F401
         Integer,  # noqa: F401
         UInt,  # noqa: F401
+        UInt2,  # noqa: F401
         UInt4,  # noqa: F401
         UInt8,  # noqa: F401
         UInt16,  # noqa: F401
         UInt32,  # noqa: F401
         UInt64,  # noqa: F401
         Int,  # noqa: F401
+        Int2,  # noqa: F401
         Int4,  # noqa: F401
         Int8,  # noqa: F401
         Int16,  # noqa: F401

test/test_array.py

@@ -78,6 +78,7 @@ def test_dtypes():
         Float64,
         Inexact,
         Int,
+        Int2,
         Int4,
         Int8,
         Int16,
@@ -86,6 +87,7 @@ def test_dtypes():
         Num,
         Shaped,
         UInt,
+        UInt2,
         UInt4,
         UInt8,
         UInt16,
@@ -156,8 +158,10 @@ def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
 
     g(jr.normal(getkey(), (3, 4)))
     g(jnp.array([[True, False]]))
+    g(jnp.array([[1, 2], [3, 4]], dtype=jnp.int2))
     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.uint2))
     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.complex64))