diff --git a/psydac/linalg/block.py b/psydac/linalg/block.py
index 02552e4cc..e8a9c44d6 100644
--- a/psydac/linalg/block.py
+++ b/psydac/linalg/block.py
@@ -804,7 +804,11 @@ def diagonal(self, *, inverse = False, out = None):
             assert isinstance(out, BlockLinearOperator)
             assert out.domain is V
             assert out.codomain is W
-            assert all(i==j for i, j in out._blocks.keys())  # is this really needed?
+
+            # Set any off-diagonal blocks to zero
+            for i, j in out.nonzero_block_indices:
+                if i != j:
+                    out[i, j] = None
         else:
             out = BlockLinearOperator(V, W)