diff --git a/docs/api/array.md b/docs/api/array.md index 445f236..6f33c9b 100644 --- a/docs/api/array.md +++ b/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 diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py index 0c9574d..7bce648 100644 --- a/jaxtyping/__init__.py +++ b/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, diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py index 9f18045..d1049af 100644 --- a/jaxtyping/_array_types.py +++ b/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] diff --git a/jaxtyping/_indirection.py b/jaxtyping/_indirection.py index 2504347..ab70b83 100644 --- a/jaxtyping/_indirection.py +++ b/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 diff --git a/test/test_array.py b/test/test_array.py index de660cf..2e99e8c 100644 --- a/test/test_array.py +++ b/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))