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

Compatibility with JAX changes to opaque dtypes

Browse Source
This commit is contained in:
Patrick Kidger
2023-07-26 09:33:32 -07:00
parent 4c90808401
commit 13e6870fb8
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
jaxtyping/_array_types.py
+10 -1
View File
@@ -137,12 +137,21 @@ def _check_dims(
return True
def _is_jax_extended_dtype(dtype: Any) -> bool:
if not has_jax:
return False
if hasattr(jax.dtypes, "extended"): # jax>=0.4.14
return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
else: # jax<=0.4.13
return jax.core.is_opaque_dtype(dtype)
class _MetaAbstractArray(type):
def __instancecheck__(cls, obj):
if not isinstance(obj, cls.array_type):
return False
if has_jax and jax.core.is_opaque_dtype(obj.dtype):
if _is_jax_extended_dtype(obj.dtype):
dtype = str(obj.dtype)
elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
# JAX, numpy