From d55a3103b0721e9b791765b7b4551dfe6bdb53bf Mon Sep 17 00:00:00 2001
From: Evgenii Zheltonozhskii <zheltonozhskiy@gmail.com>
Date: Fri, 9 Aug 2024 18:35:19 +0300
Subject: [PATCH] Fix adjoint test

---
 tests/test_adjoint.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/test_adjoint.py b/tests/test_adjoint.py
index 71c0604..7215418 100644
--- 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