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

Rewrote syntax

Browse Source
This commit is contained in:
Patrick Kidger
2022-08-30 12:47:45 -07:00
parent 3f9fad59bd
commit 903000f3d5
10 changed files with 253 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_array.py
+32 -29
View File
@@ -21,15 +21,18 @@ import jax.numpy as jnp
import jax.random as jr
import pytest
from jaxtyping import Array, f, f32, jaxtyped
from jaxtyping import f32, Float, jaxtyped, Shaped
from .helpers import ParamError, ReturnError
Array = jnp.ndarray
def test_basic(typecheck):
@jaxtyped
@typecheck
def g(x: Array["..."]):
def g(x: Shaped[Array, "..."]):
pass
g(jnp.array(1.0))
@@ -38,14 +41,14 @@ def test_basic(typecheck):
def test_return(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f["b c"]) -> f["c b"]:
def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]:
return jnp.transpose(x)
g(jr.normal(getkey(), (3, 4)))
@jaxtyped
@typecheck
def h(x: f["b c"]) -> f["b c"]:
def h(x: Float[Array, "b c"]) -> Float[Array, "b c"]:
return jnp.transpose(x)
with pytest.raises(ReturnError):
@@ -55,7 +58,7 @@ def test_return(typecheck, getkey):
def test_two_args(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: Array["b c"], y: Array["c d"]):
def g(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]):
return x @ y
g(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
@@ -64,7 +67,7 @@ def test_two_args(typecheck, getkey):
@jaxtyped
@typecheck
def h(x: Array["b c"], y: Array["c d"]) -> Array["b d"]:
def h(x: Shaped[Array, "b c"], y: Shaped[Array, "c d"]) -> Shaped[Array, "b d"]:
return x @ y
h(jr.normal(getkey(), (3, 4)), jr.normal(getkey(), (4, 5)))
@@ -75,7 +78,7 @@ def test_two_args(typecheck, getkey):
def test_any_dtype(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: Array["a b"]) -> Array["a b"]:
def g(x: Shaped[Array, "a b"]) -> Shaped[Array, "a b"]:
return x
g(jr.normal(getkey(), (3, 4)))
@@ -92,12 +95,12 @@ def test_any_dtype(typecheck, getkey):
def test_nested_jaxtyped(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["b c"], transpose: bool) -> f32["c b"]:
def g(x: f32[Array, "b c"], transpose: bool) -> f32[Array, "c b"]:
return h(x, transpose)
@jaxtyped
@typecheck
def h(x: f32["c b"], transpose: bool) -> f32["b c"]:
def h(x: f32[Array, "c b"], transpose: bool) -> f32[Array, "b c"]:
if transpose:
return jnp.transpose(x)
else:
@@ -113,11 +116,11 @@ def test_nested_jaxtyped(typecheck, getkey):
def test_nested_nojaxtyped(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["b c"]):
def g(x: f32[Array, "b c"]):
return h(x)
@typecheck
def h(x: f32["c b"]):
def h(x: f32[Array, "c b"]):
return x
with pytest.raises(ParamError):
@@ -127,14 +130,14 @@ def test_nested_nojaxtyped(typecheck, getkey):
def test_isinstance(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["b c"]) -> f32[" z"]:
def g(x: f32[Array, "b c"]) -> f32[Array, " z"]:
y = jnp.transpose(x)
assert isinstance(y, f32["c b"])
assert isinstance(y, f32[Array, "c b"])
assert not isinstance(
y, f32["b z"]
y, f32[Array, "b z"]
) # z left unbound as b!=c (unless x symmetric, which it isn't)
out = jr.normal(getkey(), (500,))
assert isinstance(out, f32["z"]) # z now bound
assert isinstance(out, f32[Array, "z"]) # z now bound
return out
g(jr.normal(getkey(), (2, 3)))
@@ -143,7 +146,7 @@ def test_isinstance(typecheck, getkey):
def test_fixed(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["4 5 foo"], y: f32[" foo"]) -> f32["4 5"]:
def g(x: f32[Array, "4 5 foo"], y: f32[Array, " foo"]) -> f32[Array, "4 5"]:
return x @ y
a = jr.normal(getkey(), (4, 5, 2))
@@ -158,7 +161,7 @@ def test_fixed(typecheck, getkey):
def test_anonymous(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["foo _"], y: f32[" _"]):
def g(x: f32[Array, "foo _"], y: f32[Array, " _"]):
pass
a = jr.normal(getkey(), (3, 4))
@@ -169,7 +172,7 @@ def test_anonymous(typecheck, getkey):
def test_named_variadic(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["*batch foo"], y: f32[" *batch"], z: f32[" foo"]):
def g(x: f32[Array, "*batch foo"], y: f32[Array, " *batch"], z: f32[Array, " foo"]):
pass
c = jr.normal(getkey(), (5,))
@@ -189,7 +192,7 @@ def test_named_variadic(typecheck, getkey):
@jaxtyped
@typecheck
def h(x: f32[" foo *batch"], y: f32[" foo *batch bar"]):
def h(x: f32[Array, " foo *batch"], y: f32[Array, " foo *batch bar"]):
pass
a = jr.normal(getkey(), (4,))
@@ -205,7 +208,7 @@ def test_named_variadic(typecheck, getkey):
def test_anonymous_variadic(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["... foo"], y: f32[" foo"]):
def g(x: f32[Array, "... foo"], y: f32[Array, " foo"]):
pass
a1 = jr.normal(getkey(), (5,))
@@ -227,7 +230,7 @@ def test_anonymous_variadic(typecheck, getkey):
def test_broadcast_fixed(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32["#4"]):
def g(x: f32[Array, "#4"]):
pass
g(jr.normal(getkey(), (4,)))
@@ -240,7 +243,7 @@ def test_broadcast_fixed(typecheck, getkey):
def test_broadcast_named(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32[" #foo"], y: f32[" #foo"]):
def g(x: f32[Array, " #foo"], y: f32[Array, " #foo"]):
pass
a = jr.normal(getkey(), (3,))
@@ -264,7 +267,7 @@ def test_broadcast_named(typecheck, getkey):
def test_broadcast_variadic_named(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: f32[" *#foo"], y: f32[" *#foo"]):
def g(x: f32[Array, " *#foo"], y: f32[Array, " *#foo"]):
pass
a = jr.normal(getkey(), (3,))
@@ -322,28 +325,28 @@ def test_broadcast_variadic_named(typecheck, getkey):
def test_no_commas():
with pytest.raises(ValueError):
f32["foo, bar"]
f32[Array, "foo, bar"]
def test_symbolic(typecheck, getkey):
@jaxtyped
@typecheck
def make_slice(x: f32[" dim"]) -> f32[" dim-1"]:
def make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]:
return x[1:]
@jaxtyped
@typecheck
def cat(x: f32[" dim"]) -> f32[" 2*dim"]:
def cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]:
return jnp.concatenate([x, x])
@jaxtyped
@typecheck
def bad_make_slice(x: f32[" dim"]) -> f32[" dim-1"]:
def bad_make_slice(x: f32[Array, " dim"]) -> f32[Array, " dim-1"]:
return x
@jaxtyped
@typecheck
def bad_cat(x: f32[" dim"]) -> f32[" 2*dim"]:
def bad_cat(x: f32[Array, " dim"]) -> f32[Array, " 2*dim"]:
return jnp.concatenate([x, x, x])
x = jr.normal(getkey(), (5,))
@@ -365,7 +368,7 @@ def test_symbolic(typecheck, getkey):
def test_incomplete_symbolic(typecheck, getkey):
@jaxtyped
@typecheck
def foo(x: f32[" 2*dim"]):
def foo(x: f32[Array, " 2*dim"]):
pass
x = jr.normal(getkey(), (4,))