wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* Add int4/uint4 support in jaxtyping.

* Fix a typo and update api docs.
This commit is contained in:
jianlijianli
2024-02-25 12:07:01 +00:00
committed by Patrick Kidger
parent 1d4d40294c
commit 9beb5f2d29
5 changed files with 18 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_array.py
+4
View File
@@ -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))