From da8300ec6cc6d7f9a66c96009b7f7399f2bd8208 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Tue, 20 Sep 2022 15:16:52 -0700
Subject: [PATCH] tidyness tweak (#28)

---
 jaxtyping/array_types.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py
index 3a91678..2e18b1d 100644
--- a/jaxtyping/array_types.py
+++ b/jaxtyping/array_types.py
@@ -150,26 +150,25 @@ class _MetaAbstractArray(type):
         if cls.dtypes is not _any_dtype and dtype not in cls.dtypes:
             return False
 
-        temp_memo = not hasattr(storage, "memo_stack") or len(storage.memo_stack) == 0
+        no_temp_memo = hasattr(storage, "memo_stack") and len(storage.memo_stack) != 0
 
-        if temp_memo:
+        if no_temp_memo:
+            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
+            # Make a copy so we don't mutate the original memo during the shape check.
+            single_memo = single_memo.copy()
+            variadic_memo = variadic_memo.copy()
+            variadic_broadcast_memo = variadic_broadcast_memo.copy()
+        else:
             # `isinstance` happening outside any @jaxtyped decorators, e.g. at the
             # global scope. In this case just create a temporary memo, since we're not
             # going to be comparing against any stored values anyway.
             single_memo = {}
             variadic_memo = {}
             variadic_broadcast_memo = {}
-        else:
-            single_memo, variadic_memo, variadic_broadcast_memo = storage.memo_stack[-1]
-            # Make a copy so we don't mutate the original memo during the shape check.
-            single_memo = single_memo.copy()
-            variadic_memo = variadic_memo.copy()
-            variadic_broadcast_memo = variadic_broadcast_memo.copy()
-            temp_memo = False
 
         if cls._check_shape(obj, single_memo, variadic_memo, variadic_broadcast_memo):
             # We update the memo every time we successfully pass a shape check
-            if not temp_memo:
+            if no_temp_memo:
                 storage.memo_stack[-1] = (
                     single_memo,
                     variadic_memo,