Fix adjoint test

patrick-kidger · Aug 9, 2024 · d55a310 · d55a310
1 parent c3af158
commit d55a310
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py
@@ -7,6 +7,7 @@
 
 from .helpers import (
     make_diagonal_operator,
+    make_identity_operator,
     make_operators,
     make_tridiagonal_operator,
     tree_allclose,
@@ -16,7 +17,10 @@
 @pytest.mark.parametrize("make_operator", make_operators)
 @pytest.mark.parametrize("dtype", (jnp.float64, jnp.complex128))
 def test_adjoint(make_operator, dtype, getkey):
-    if make_operator is make_diagonal_operator:
+    if (
+        make_operator is make_diagonal_operator
+        or make_operator is make_identity_operator
+    ):
         matrix = jnp.eye(4, dtype=dtype)
         tags = lx.diagonal_tag
         in_size = out_size = 4