Add support for numpy struct dtype

jaxtyping/__init__.py

@@ -28,6 +28,7 @@
  AbstractArray as AbstractArray,
  AbstractDtype as AbstractDtype,
  get_array_name_format as get_array_name_format,
+ make_numpy_struct_dtype as make_numpy_struct_dtype,
  set_array_name_format as set_array_name_format,
 )
 from ._config import config as config

jaxtyping/_array_types.py

@@ -157,6 +157,10 @@ def _check_dims(
  return ""
 
 
+def _dtype_is_numpy_struct_array(dtype):
+ return dtype.type.__name__ == "void" and dtype is not np.dtype(np.void)
+
+
 class _MetaAbstractArray(type):
  _skip_instancecheck: bool = False
 
@@ -177,6 +181,9 @@ def __instancecheck_str__(cls, obj: Any) -> str:
  if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
  # JAX, numpy
  dtype = obj.dtype.type.__name__
+ # numpy structured array is strictly a subtype of np.void
+ if _dtype_is_numpy_struct_array(obj.dtype):
+ dtype = str(obj.dtype)
  elif hasattr(obj.dtype, "as_numpy_dtype"):
  # TensorFlow
  dtype = obj.dtype.as_numpy_dtype.__name__
@@ -755,3 +762,34 @@ class _Cls(AbstractDtype):
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 Key = _make_dtype(_prng_key, "Key")
+
+
+def make_numpy_struct_dtype(dtype: np.dtype, name: str):
+ """Creates a type annotation for [numpy structured array](https://numpy.org/doc/stable/user/basics.rec.html#structured-arrays)
+ It does exact match on the name, order, and dtype of all its fields.
+
+ !!! Example
+
+ ```python
+ label_t = np.dtype([('first', np.uint8), ('second', np.int8)])
+ Label = make_numpy_struct_dtype(label_t, 'Label')
+ ```
+ after that, you can use it just like any AbstractDtype
+ ```python
+ a: Label[np.ndarray, 'a b'] = np.array([[(1, 0), (0, 1)]], dtype=label_t)
+ ```
+
+ **Arguments:**
+
+ - `dtype`: The numpy dtype that the returned annotation matches
+
+ - `name`: The python class name for the returned dtype annotation
+
+ **Returns:**
+
+ A type annotation with classname `name` and matching exactly `dtype`.
+ It can be used like any usual subclasses of AbstractDtypes.
+ """
+ if not (isinstance(dtype, np.dtype) and _dtype_is_numpy_struct_array(dtype)):
+ raise ValueError(f"Expecting a numpy structured array dtype, not {dtype}")
+ return _make_dtype(str(dtype), name)

test/test_array.py

@@ -92,6 +92,24 @@ def test_dtypes():
  assert key == val.__name__
 
 
+def test_numpy_struct_dtype():
+ from jaxtyping import make_numpy_struct_dtype
+
+ dtype1 = np.dtype([("first", np.uint8), ("second", bool)])
+ Dtype1 = make_numpy_struct_dtype(dtype1, "Dtype1")
+ arr = np.array([0, False], dtype=dtype1)
+
+ assert isinstance(arr, Dtype1[np.ndarray, "_"])
+
+ dtype2 = np.dtype([("third", np.uint8), ("second", bool)])
+ Dtype2 = make_numpy_struct_dtype(dtype2, "Dtype2")
+ assert not isinstance(arr, Dtype2[np.ndarray, "_"])
+
+ dtype3 = np.dtype([("second", bool), ("first", np.uint8)])
+ Dtype3 = make_numpy_struct_dtype(dtype3, "Dtype3")
+ assert not isinstance(arr, Dtype3[np.ndarray, "_"])
+
+
 def test_return(jaxtyp, typecheck, getkey):
  @jaxtyp(typecheck)
  def g(x: Float[Array, "b c"]) -> Float[Array, "c b"]: