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

Replace ndim with len(shape)

Browse Source
This commit is contained in:
Piotr Kaminski
2024-03-11 20:35:29 +01:00
committed by Patrick Kidger
parent 696cc5b90f
commit 07e58de0bd
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/_array_types.py
+4 -4
View File
@@ -224,12 +224,12 @@ class _MetaAbstractArray(type):
arg_memo: dict[str, Any],
) -> str:
if cls.index_variadic is None:
if obj.ndim != len(cls.dims):
return f"this array has {obj.ndim} dimensions, not the {len(cls.dims)} expected by the type hint" # noqa: E501
if len(obj.shape) != len(cls.dims):
return f"this array has {len(obj.shape)} dimensions, not the {len(cls.dims)} expected by the type hint" # noqa: E501
return _check_dims(cls.dims, obj.shape, single_memo, arg_memo)
else:
if obj.ndim < len(cls.dims) - 1:
return f"this array has {obj.ndim} dimensions, which is fewer than {len(cls.dims) - 1} that is the minimum expected by the type hint" # noqa: E501
if len(obj.shape) < len(cls.dims) - 1:
return f"this array has {len(obj.shape)} dimensions, which is fewer than {len(cls.dims) - 1} that is the minimum expected by the type hint" # noqa: E501
i = cls.index_variadic
j = -(len(cls.dims) - i - 1)
if j == 0: