diff --git a/setup.py b/setup.py
index d9831e15..f82ce6d0 100644
--- a/setup.py
+++ b/setup.py
@@ -53,6 +53,7 @@
     },
     extras_require={
         "test": [
+            "autoray<0.7.0", # autoray.tensorflow_diag no longer works
             "docutils>=0.19",
             "flaky",
             "pre-commit",