warp-lang 1.3.0__py3-none-macosx_10_13_universal2.whl → 1.3.2__py3-none-macosx_10_13_universal2.whl

warp/stubs.py

@@ -477,7 +477,7 @@ def cross(a: Vector[3, Scalar], b: Vector[3, Scalar]) -> Vector[3, Scalar]:
 
 
 @over
-def skew(vec: Vector[3, Scalar]):
+def skew(vec: Vector[3, Scalar]) -> Matrix[3, 3, Scalar]:
     """Compute the skew-symmetric 3x3 matrix for a 3D vector ``vec``."""
     ...
 
@@ -519,7 +519,7 @@ def normalize(a: Quaternion[Float]) -> Quaternion[Float]:
 
 
 @over
-def transpose(a: Matrix[Any, Any, Scalar]):
+def transpose(a: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
     """Return the transpose of the matrix ``a``."""
     ...
 
@@ -747,13 +747,13 @@ def spatial_cross_dual(a: Vector[6, Float], b: Vector[6, Float]) -> Vector[6, Fl
 
 
 @over
-def spatial_top(svec: Vector[6, Float]):
+def spatial_top(svec: Vector[6, Float]) -> Vector[3, Float]:
     """Return the top (first) part of a 6D screw vector."""
     ...
 
 
 @over
-def spatial_bottom(svec: Vector[6, Float]):
+def spatial_bottom(svec: Vector[6, Float]) -> Vector[3, Float]:
     """Return the bottom (second) part of a 6D screw vector."""
     ...
 
@@ -1140,212 +1140,212 @@ def tid() -> Tuple[int, int, int, int]:
 
 
 @over
-def select(cond: bool, value_if_false: Any, value_if_true: Any):
+def select(cond: bool, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: int8, value_if_false: Any, value_if_true: Any):
+def select(cond: int8, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: uint8, value_if_false: Any, value_if_true: Any):
+def select(cond: uint8, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: int16, value_if_false: Any, value_if_true: Any):
+def select(cond: int16, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: uint16, value_if_false: Any, value_if_true: Any):
+def select(cond: uint16, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: int32, value_if_false: Any, value_if_true: Any):
+def select(cond: int32, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: uint32, value_if_false: Any, value_if_true: Any):
+def select(cond: uint32, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: int64, value_if_false: Any, value_if_true: Any):
+def select(cond: int64, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(cond: uint64, value_if_false: Any, value_if_true: Any):
+def select(cond: uint64, value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def select(arr: Array[Any], value_if_false: Any, value_if_true: Any):
+def select(arr: Array[Any], value_if_false: Any, value_if_true: Any) -> Any:
     """Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``"""
     ...
 
 
 @over
-def atomic_add(arr: Array[Any], i: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i]``."""
+def atomic_add(arr: Array[Any], i: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: Array[Any], i: int32, j: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j]``."""
+def atomic_add(arr: Array[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k]``."""
+def atomic_add(arr: Array[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_add(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: FabricArray[Any], i: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i]``."""
+def atomic_add(arr: FabricArray[Any], i: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: FabricArray[Any], i: int32, j: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j]``."""
+def atomic_add(arr: FabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k]``."""
+def atomic_add(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_add(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: IndexedFabricArray[Any], i: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i]``."""
+def atomic_add(arr: IndexedFabricArray[Any], i: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j]``."""
+def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k]``."""
+def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically add ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_add(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: Array[Any], i: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i]``."""
+def atomic_sub(arr: Array[Any], i: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: Array[Any], i: int32, j: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j]``."""
+def atomic_sub(arr: Array[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k]``."""
+def atomic_sub(arr: Array[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_sub(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: FabricArray[Any], i: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i]``."""
+def atomic_sub(arr: FabricArray[Any], i: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j]``."""
+def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k]``."""
+def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_sub(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: IndexedFabricArray[Any], i: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i]``."""
+def atomic_sub(arr: IndexedFabricArray[Any], i: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j]``."""
+def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k]``."""
+def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
     ...
 
 
 @over
-def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]``."""
+def atomic_sub(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
     ...
 
 
 @over
-def atomic_min(arr: Array[Any], i: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_min(arr: Array[Any], i: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1353,8 +1353,8 @@ def atomic_min(arr: Array[Any], i: int32, value: Any):
 
 
 @over
-def atomic_min(arr: Array[Any], i: int32, j: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_min(arr: Array[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1362,8 +1362,8 @@ def atomic_min(arr: Array[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1371,8 +1371,8 @@ def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
 
 
 @over
-def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1380,8 +1380,8 @@ def atomic_min(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: A
 
 
 @over
-def atomic_min(arr: FabricArray[Any], i: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_min(arr: FabricArray[Any], i: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1389,8 +1389,8 @@ def atomic_min(arr: FabricArray[Any], i: int32, value: Any):
 
 
 @over
-def atomic_min(arr: FabricArray[Any], i: int32, j: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_min(arr: FabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1398,8 +1398,8 @@ def atomic_min(arr: FabricArray[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1407,8 +1407,8 @@ def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
 
 
 @over
-def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1416,8 +1416,8 @@ def atomic_min(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, va
 
 
 @over
-def atomic_min(arr: IndexedFabricArray[Any], i: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_min(arr: IndexedFabricArray[Any], i: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1425,8 +1425,8 @@ def atomic_min(arr: IndexedFabricArray[Any], i: int32, value: Any):
 
 
 @over
-def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1434,8 +1434,8 @@ def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1443,8 +1443,8 @@ def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value
 
 
 @over
-def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the minimum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1452,8 +1452,8 @@ def atomic_min(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: in
 
 
 @over
-def atomic_max(arr: Array[Any], i: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_max(arr: Array[Any], i: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1461,8 +1461,8 @@ def atomic_max(arr: Array[Any], i: int32, value: Any):
 
 
 @over
-def atomic_max(arr: Array[Any], i: int32, j: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_max(arr: Array[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1470,8 +1470,8 @@ def atomic_max(arr: Array[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1479,8 +1479,8 @@ def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, value: Any):
 
 
 @over
-def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1488,8 +1488,8 @@ def atomic_max(arr: Array[Any], i: int32, j: int32, k: int32, l: int32, value: A
 
 
 @over
-def atomic_max(arr: FabricArray[Any], i: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_max(arr: FabricArray[Any], i: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1497,8 +1497,8 @@ def atomic_max(arr: FabricArray[Any], i: int32, value: Any):
 
 
 @over
-def atomic_max(arr: FabricArray[Any], i: int32, j: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_max(arr: FabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1506,8 +1506,8 @@ def atomic_max(arr: FabricArray[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1515,8 +1515,8 @@ def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, value: Any):
 
 
 @over
-def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1524,8 +1524,8 @@ def atomic_max(arr: FabricArray[Any], i: int32, j: int32, k: int32, l: int32, va
 
 
 @over
-def atomic_max(arr: IndexedFabricArray[Any], i: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i]`` and atomically update the array.
+def atomic_max(arr: IndexedFabricArray[Any], i: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1533,8 +1533,8 @@ def atomic_max(arr: IndexedFabricArray[Any], i: int32, value: Any):
 
 
 @over
-def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j]`` and atomically update the array.
+def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1542,8 +1542,8 @@ def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, value: Any):
 
 
 @over
-def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1551,8 +1551,8 @@ def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, value
 
 
 @over
-def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any):
-    """Compute the maximum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+def atomic_max(arr: IndexedFabricArray[Any], i: int32, j: int32, k: int32, l: int32, value: Any) -> Any:
+    """Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
     .. note:: The operation is only atomic on a per-component basis for vectors and matrices.
     """
@@ -1778,7 +1778,7 @@ def mul(a: Vector[Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Vector[Any, Scal
 
 
 @over
-def mul(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]):
+def mul(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
     """ """
     ...
 
@@ -1803,7 +1803,7 @@ def mul(a: Transformation[Scalar], b: Scalar) -> Transformation[Scalar]:
 
 @over
 def mod(a: Scalar, b: Scalar) -> Scalar:
-    """ """
+    """Modulo operation using truncated division."""
     ...
 
 

warp/tape.py

@@ -197,7 +197,7 @@ class Tape:
         else:
             self.scopes.append((len(self.launches), None, None))
 
-    def check_kernel_array_access(self, kernel, args):
+    def _check_kernel_array_access(self, kernel, args):
         """Detect illegal inter-kernel write after read access patterns during launch capture"""
         adj = kernel.adj
         kernel_name = adj.fun_name
@@ -256,7 +256,7 @@ class Tape:
         self.scopes = []
         self.zero()
         if wp.config.verify_autograd_array_access:
-            self.reset_array_read_flags()
+            self._reset_array_read_flags()
 
     def zero(self):
         """
@@ -271,7 +271,7 @@ class Tape:
                 else:
                     g.zero_()
 
-    def reset_array_read_flags(self):
+    def _reset_array_read_flags(self):
         """
         Reset all recorded array read flags to False
         """

warp/tests/test_array.py

@@ -2382,6 +2382,14 @@ def test_kernel_array_from_ptr(test, device):
     assert_np_equal(arr.numpy(), np.array(((1.0, 2.0, 3.0), (0.0, 0.0, 0.0))))
 
 
+def test_array_from_int32_domain(test, device):
+    wp.zeros(np.array([1504, 1080, 520], dtype=np.int32), dtype=wp.float32, device=device)
+
+
+def test_array_from_int64_domain(test, device):
+    wp.zeros(np.array([1504, 1080, 520], dtype=np.int64), dtype=wp.float32, device=device)
+
+
 devices = get_test_devices()
 
 
@@ -2443,6 +2451,9 @@ add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, dev
 add_function_test(TestArray, "test_direct_from_numpy", test_direct_from_numpy, devices=["cpu"])
 add_function_test(TestArray, "test_kernel_array_from_ptr", test_kernel_array_from_ptr, devices=devices)
 
+add_function_test(TestArray, "test_array_from_int32_domain", test_array_from_int32_domain, devices=devices)
+add_function_test(TestArray, "test_array_from_int64_domain", test_array_from_int64_domain, devices=devices)
+
 try:
     import torch
 

warp/tests/test_async.py

@@ -21,7 +21,9 @@ class Capturable:
 
     def __enter__(self):
         if self.use_graph:
-            wp.capture_begin(stream=self.stream)
+            # preload module before graph capture
+            wp.load_module(device=wp.get_device())
+            wp.capture_begin(stream=self.stream, force_module_load=False)
 
     def __exit__(self, exc_type, exc_value, traceback):
         if self.use_graph: