wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-11 12:21:38 +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_pytree.py
+3 -3
View File
@@ -25,7 +25,7 @@ import jax.numpy as jnp
import jax.random as jr
import pytest
from jaxtyping import f, jaxtyped, PyTree
from jaxtyping import Float, jaxtyped, PyTree
from .helpers import make_mlp, ParamError
@@ -95,7 +95,7 @@ def test_nested_pytrees(getkey, typecheck):
def test_pytree_array(typecheck):
@jaxtyped
@typecheck
def g(x: PyTree[f["..."]]):
def g(x: PyTree[Float[jnp.ndarray, "..."]]):
pass
g(jnp.array(1.0))
@@ -109,7 +109,7 @@ def test_pytree_array(typecheck):
def test_pytree_shaped_array(typecheck, getkey):
@jaxtyped
@typecheck
def g(x: PyTree[f["b c"]]):
def g(x: PyTree[Float[jnp.ndarray, "b c"]]):
pass
g(jnp.array([[1.0]]))