merge resolve

patrick-kidger · Mar 4, 2024 · 0180958 · 0180958
1 parent 83b376b
commit 0180958
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_ad.py b/tests/test_ad.py
@@ -480,7 +480,7 @@ def f_bwd(res, g):
     f_custom.defvjp(f_fwd, f_bwd)
     jax_hess = jax.hessian(f)(1.0, 2.0)
     eqx_hess = eqx.filter_hessian(f)(jnp.array(1.0), jnp.array(2.0))
-    assert shaped_allclose(jax_hess, eqx_hess)
+    assert tree_allclose(jax_hess, eqx_hess)
 
 
 def test_filter_custom_jvp_symbolic_zero():