Extra test

danielward27 · Sep 2, 2024 · e53c865 · e53c865
1 parent 15238b9
commit e53c865
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
@@ -37,6 +37,14 @@ def test_Parameterize():
     assert pytest.approx(unwrap(unwrappable)) == jnp.zeros((3, 2))
 
 
+def test_nested_Parameterized():
+    param = Parameterize(
+        jnp.square,
+        Parameterize(jnp.square, Parameterize(jnp.square, 2)),
+    )
+    assert unwrap(param) == jnp.square(jnp.square(jnp.square(2)))
+
+
 def test_NonTrainable_and_non_trainable():
     dist1 = eqx.tree_at(lambda dist: dist.bijection, Normal(), replace_fn=NonTrainable)
     dist2 = non_trainable(Normal())