Add IntScalar type

probml · Oct 16, 2024 · 42c5c22 · 42c5c22
1 parent 75205f8
commit 42c5c22
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/dynamax/types.py b/dynamax/types.py
@@ -1,6 +1,8 @@
 from typing import Union
-from jaxtyping import Array, Float
+from jaxtyping import Array, Float, Int
 
 PRNGKeyT = Array
 
 Scalar = Union[float, Float[Array, ""]] # python float or scalar jax device array with dtype float
+
+IntScalar = Union[int, Int[Array, ""]]