wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

tidyness tweak (#28)

Browse Source
This commit is contained in:
Patrick Kidger
2022-09-20 15:16:52 -07:00
committed by GitHub
parent 39439c2790
commit da8300ec6c
1 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/array_types.py
+9 -10
View File
@@ -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,