From d3651ca70e3a748ec74a3300668e742dacfc25c4 Mon Sep 17 00:00:00 2001 From: Peter Roelants Date: Mon, 3 Oct 2022 16:40:51 +0200 Subject: [PATCH] NamedTuple example (#36) --- test/test_pytree.py | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/test/test_pytree.py b/test/test_pytree.py index 244ebaf..8144858 100644 --- a/test/test_pytree.py +++ b/test/test_pytree.py @@ -17,7 +17,7 @@ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -from typing import Tuple, Union +from typing import NamedTuple, Tuple, Union import equinox as eqx import jax @@ -155,3 +155,33 @@ def test_pytree_tuple(typecheck): g([1, 1]) with pytest.raises(ParamError): g([(1, 1), "hi"]) + + +def test_pytree_namedtuple(typecheck): + class CustomNamedTuple(NamedTuple): + x: Float[jnp.ndarray, "a b"] + y: Float[jnp.ndarray, "b c"] + + class OtherCustomNamedTuple(NamedTuple): + x: Float[jnp.ndarray, "a b"] + y: Float[jnp.ndarray, "b c"] + + @typecheck + def g(x: PyTree[CustomNamedTuple]): + ... + + g( + CustomNamedTuple( + x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)), + y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)), + ) + ) + with pytest.raises(ParamError): + g(object()) + with pytest.raises(ParamError): + g( + OtherCustomNamedTuple( + x=jax.random.normal(jax.random.PRNGKey(42), (3, 2)), + y=jax.random.normal(jax.random.PRNGKey(420), (2, 5)), + ) + )