warp-lang 1.1.0__py3-none-manylinux2014_aarch64.whl → 1.2.1__py3-none-manylinux2014_aarch64.whl

warp/sim/integrator_xpbd.py

@@ -972,11 +972,9 @@ def apply_body_delta_velocities(
 
 
 @wp.kernel
-def apply_joint_torques(
+def apply_joint_actions(
     body_q: wp.array(dtype=wp.transform),
     body_com: wp.array(dtype=wp.vec3),
-    joint_q_start: wp.array(dtype=int),
-    joint_qd_start: wp.array(dtype=int),
     joint_type: wp.array(dtype=int),
     joint_parent: wp.array(dtype=int),
     joint_child: wp.array(dtype=int),
@@ -1028,8 +1026,6 @@ def apply_joint_torques(
     # q_c = wp.transform_get_rotation(X_wc)
 
     # joint properties (for 1D joints)
-    # q_start = joint_q_start[tid]
-    qd_start = joint_qd_start[tid]
     axis_start = joint_axis_start[tid]
     lin_axis_count = joint_axis_dim[tid, 0]
     ang_axis_count = joint_axis_dim[tid, 1]
@@ -1043,14 +1039,14 @@ def apply_joint_torques(
         mode = joint_axis_mode[axis_start]
         if mode == wp.sim.JOINT_MODE_FORCE:
             axis = joint_axis[axis_start]
-            act = joint_act[qd_start]
+            act = joint_act[axis_start]
             a_p = wp.transform_vector(X_wp, axis)
             t_total += act * a_p
     elif type == wp.sim.JOINT_PRISMATIC:
         mode = joint_axis_mode[axis_start]
         if mode == wp.sim.JOINT_MODE_FORCE:
             axis = joint_axis[axis_start]
-            act = joint_act[qd_start]
+            act = joint_act[axis_start]
             a_p = wp.transform_vector(X_wp, axis)
             f_total += act * a_p
     elif type == wp.sim.JOINT_COMPOUND:
@@ -1075,13 +1071,13 @@ def apply_joint_torques(
 
         if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
             axis_0 = joint_axis[axis_start + 0]
-            t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0)
+            t_total += joint_act[axis_start + 0] * wp.transform_vector(X_wp, axis_0)
         if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
             axis_1 = joint_axis[axis_start + 1]
-            t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1)
+            t_total += joint_act[axis_start + 1] * wp.transform_vector(X_wp, axis_1)
         if joint_axis_mode[axis_start + 2] == wp.sim.JOINT_MODE_FORCE:
             axis_2 = joint_axis[axis_start + 2]
-            t_total += joint_act[qd_start + 2] * wp.transform_vector(X_wp, axis_2)
+            t_total += joint_act[axis_start + 2] * wp.transform_vector(X_wp, axis_2)
 
     elif type == wp.sim.JOINT_UNIVERSAL:
         # q_off = wp.transform_get_rotation(X_cj)
@@ -1107,10 +1103,10 @@ def apply_joint_torques(
 
         if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
             axis_0 = joint_axis[axis_start + 0]
-            t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0)
+            t_total += joint_act[axis_start + 0] * wp.transform_vector(X_wp, axis_0)
         if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
             axis_1 = joint_axis[axis_start + 1]
-            t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1)
+            t_total += joint_act[axis_start + 1] * wp.transform_vector(X_wp, axis_1)
 
     elif type == wp.sim.JOINT_D6:
         # unroll for loop to ensure joint actions remain differentiable
@@ -1119,43 +1115,43 @@ def apply_joint_torques(
         if lin_axis_count > 0:
             if joint_axis_mode[axis_start + 0] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + 0]
-                act = joint_act[qd_start + 0]
+                act = joint_act[axis_start + 0]
                 a_p = wp.transform_vector(X_wp, axis)
                 f_total += act * a_p
         if lin_axis_count > 1:
             if joint_axis_mode[axis_start + 1] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + 1]
-                act = joint_act[qd_start + 1]
+                act = joint_act[axis_start + 1]
                 a_p = wp.transform_vector(X_wp, axis)
                 f_total += act * a_p
         if lin_axis_count > 2:
             if joint_axis_mode[axis_start + 2] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + 2]
-                act = joint_act[qd_start + 2]
+                act = joint_act[axis_start + 2]
                 a_p = wp.transform_vector(X_wp, axis)
                 f_total += act * a_p
 
         if ang_axis_count > 0:
             if joint_axis_mode[axis_start + lin_axis_count + 0] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + lin_axis_count + 0]
-                act = joint_act[qd_start + lin_axis_count + 0]
+                act = joint_act[axis_start + lin_axis_count + 0]
                 a_p = wp.transform_vector(X_wp, axis)
                 t_total += act * a_p
         if ang_axis_count > 1:
             if joint_axis_mode[axis_start + lin_axis_count + 1] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + lin_axis_count + 1]
-                act = joint_act[qd_start + lin_axis_count + 1]
+                act = joint_act[axis_start + lin_axis_count + 1]
                 a_p = wp.transform_vector(X_wp, axis)
                 t_total += act * a_p
         if ang_axis_count > 2:
             if joint_axis_mode[axis_start + lin_axis_count + 2] == wp.sim.JOINT_MODE_FORCE:
                 axis = joint_axis[axis_start + lin_axis_count + 2]
-                act = joint_act[qd_start + lin_axis_count + 2]
+                act = joint_act[axis_start + lin_axis_count + 2]
                 a_p = wp.transform_vector(X_wp, axis)
                 t_total += act * a_p
 
     else:
-        print("joint type not handled in apply_joint_torques")
+        print("joint type not handled in apply_joint_actions")
 
     # write forces
     if id_p >= 0:
@@ -2838,13 +2834,11 @@ class XPBDIntegrator(Integrator):
 
                 if model.joint_count:
                     wp.launch(
-                        kernel=apply_joint_torques,
+                        kernel=apply_joint_actions,
                         dim=model.joint_count,
                         inputs=[
                             state_in.body_q,
                             model.body_com,
-                            model.joint_q_start,
-                            model.joint_qd_start,
                             model.joint_type,
                             model.joint_parent,
                             model.joint_child,

warp/sparse.py

@@ -128,8 +128,8 @@ def bsr_zeros(
 
     bsr = bsr_matrix_t(block_type)()
 
-    bsr.nrow = rows_of_blocks
-    bsr.ncol = cols_of_blocks
+    bsr.nrow = int(rows_of_blocks)
+    bsr.ncol = int(cols_of_blocks)
     bsr.nnz = 0
     bsr.columns = wp.empty(shape=(0,), dtype=int, device=device)
     bsr.values = wp.empty(shape=(0,), dtype=block_type, device=device)
@@ -152,7 +152,11 @@ def _bsr_ensure_fits(bsr: BsrMatrix, nrow: int = None, nnz: int = None):
         bsr.values = wp.empty(shape=(nnz,), dtype=bsr.values.dtype, device=bsr.values.device)
 
 
-def bsr_set_zero(bsr: BsrMatrix, rows_of_blocks: Optional[int] = None, cols_of_blocks: Optional[int] = None):
+def bsr_set_zero(
+    bsr: BsrMatrix,
+    rows_of_blocks: Optional[int] = None,
+    cols_of_blocks: Optional[int] = None,
+):
     """
     Sets a BSR matrix to zero, possibly changing its size
 
@@ -163,9 +167,9 @@ def bsr_set_zero(bsr: BsrMatrix, rows_of_blocks: Optional[int] = None, cols_of_b
     """
 
     if rows_of_blocks is not None:
-        bsr.nrow = rows_of_blocks
+        bsr.nrow = int(rows_of_blocks)
     if cols_of_blocks is not None:
-        bsr.ncol = cols_of_blocks
+        bsr.ncol = int(cols_of_blocks)
     bsr.nnz = 0
     _bsr_ensure_fits(bsr)
     bsr.offsets.zero_()
@@ -254,7 +258,10 @@ def bsr_set_from_triplets(
     )
 
 
-def bsr_assign(dest: BsrMatrix[BlockType[Rows, Cols, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Any]]):
+def bsr_assign(
+    dest: BsrMatrix[BlockType[Rows, Cols, Scalar]],
+    src: BsrMatrix[BlockType[Rows, Cols, Any]],
+):
     """Copies the content of the `src` matrix to `dest`, casting the block values if the two matrices use distinct scalar types."""
 
     if dest.values.device != src.values.device:
@@ -288,12 +295,20 @@ def bsr_copy(A: BsrMatrix, scalar_type: Optional[Scalar] = None):
     else:
         block_type = wp.types.matrix(shape=A.block_shape, dtype=scalar_type)
 
-    copy = bsr_zeros(rows_of_blocks=A.nrow, cols_of_blocks=A.ncol, block_type=block_type, device=A.values.device)
+    copy = bsr_zeros(
+        rows_of_blocks=A.nrow,
+        cols_of_blocks=A.ncol,
+        block_type=block_type,
+        device=A.values.device,
+    )
     bsr_assign(dest=copy, src=A)
     return copy
 
 
-def bsr_set_transpose(dest: BsrMatrix[BlockType[Cols, Rows, Scalar]], src: BsrMatrix[BlockType[Rows, Cols, Scalar]]):
+def bsr_set_transpose(
+    dest: BsrMatrix[BlockType[Cols, Rows, Scalar]],
+    src: BsrMatrix[BlockType[Rows, Cols, Scalar]],
+):
     """Assigns the transposed matrix `src` to matrix `dest`"""
 
     if dest.values.device != src.values.device:
@@ -356,7 +371,12 @@ def bsr_transposed(A: BsrMatrix):
     else:
         block_type = wp.types.matrix(shape=A.block_shape[::-1], dtype=A.scalar_type)
 
-    transposed = bsr_zeros(rows_of_blocks=A.ncol, cols_of_blocks=A.nrow, block_type=block_type, device=A.values.device)
+    transposed = bsr_zeros(
+        rows_of_blocks=A.ncol,
+        cols_of_blocks=A.nrow,
+        block_type=block_type,
+        device=A.values.device,
+    )
     bsr_set_transpose(dest=transposed, src=A)
     return transposed
 
@@ -399,7 +419,10 @@ def bsr_get_diag(A: BsrMatrix[_BlockType], out: "Optional[Array[BlockType]]" = N
             raise ValueError(f"Output array must be of length at least {dim}")
 
     wp.launch(
-        kernel=_bsr_get_diag_kernel, dim=dim, device=A.values.device, inputs=[A.offsets, A.columns, A.values, out]
+        kernel=_bsr_get_diag_kernel,
+        dim=dim,
+        device=A.values.device,
+        inputs=[A.offsets, A.columns, A.values, out],
     )
 
     return out
@@ -567,7 +590,9 @@ def bsr_set_identity(A: BsrMatrix, rows_of_blocks: Optional[int] = None):
 
 
 def bsr_identity(
-    rows_of_blocks: int, block_type: BlockType[Rows, Rows, Scalar], device: wp.context.Devicelike = None
+    rows_of_blocks: int,
+    block_type: BlockType[Rows, Rows, Scalar],
+    device: wp.context.Devicelike = None,
 ) -> BsrMatrix[BlockType[Rows, Rows, Scalar]]:
     """Creates and returns a square identity matrix.
 
@@ -576,7 +601,12 @@ def bsr_identity(
         block_type: Block type for the newly created matrix -- must be square
         device: Device onto which to allocate the data arrays
     """
-    A = bsr_zeros(rows_of_blocks=rows_of_blocks, cols_of_blocks=rows_of_blocks, block_type=block_type, device=device)
+    A = bsr_zeros(
+        rows_of_blocks=rows_of_blocks,
+        cols_of_blocks=rows_of_blocks,
+        block_type=block_type,
+        device=device,
+    )
     bsr_set_identity(A)
     return A
 
@@ -601,7 +631,12 @@ def bsr_scale(x: BsrMatrix, alpha: Scalar) -> BsrMatrix:
             if not isinstance(alpha, x.scalar_type):
                 alpha = x.scalar_type(alpha)
 
-            wp.launch(kernel=_bsr_scale_kernel, dim=x.nnz, device=x.values.device, inputs=[alpha, x.values])
+            wp.launch(
+                kernel=_bsr_scale_kernel,
+                dim=x.nnz,
+                device=x.values.device,
+                inputs=[alpha, x.values],
+            )
 
     return x
 
@@ -723,10 +758,20 @@ def bsr_axpy(
     work_arrays._allocate(device, y, sum_nnz)
 
     wp.copy(work_arrays._sum_cols, y.columns, 0, 0, y.nnz)
-    wp.launch(kernel=_bsr_get_block_row, device=device, dim=y.nnz, inputs=[0, y.offsets, work_arrays._sum_rows])
+    wp.launch(
+        kernel=_bsr_get_block_row,
+        device=device,
+        dim=y.nnz,
+        inputs=[0, y.offsets, work_arrays._sum_rows],
+    )
 
     wp.copy(work_arrays._sum_cols, x.columns, y.nnz, 0, x.nnz)
-    wp.launch(kernel=_bsr_get_block_row, device=device, dim=x.nnz, inputs=[y.nnz, x.offsets, work_arrays._sum_rows])
+    wp.launch(
+        kernel=_bsr_get_block_row,
+        device=device,
+        dim=x.nnz,
+        inputs=[y.nnz, x.offsets, work_arrays._sum_rows],
+    )
 
     # Save old y values before overwriting matrix
     wp.copy(dest=work_arrays._old_y_values, src=y.values, count=y.nnz)
@@ -858,22 +903,25 @@ def _bsr_mm_compute_values(
     mm_cols: wp.array(dtype=int),
     mm_values: wp.array(dtype=Any),
 ):
-    row = wp.tid()
-    mm_beg = mm_offsets[row]
-    mm_end = mm_offsets[row + 1]
+    mm_block = wp.tid()
+
+    row = wp.lower_bound(mm_offsets, mm_block + 1) - 1
+    col = mm_cols[mm_block]
+
+    mm_val = mm_values.dtype(type(alpha)(0.0))
 
     x_beg = x_offsets[row]
     x_end = x_offsets[row + 1]
     for x_block in range(x_beg, x_end):
         x_col = x_columns[x_block]
-        ax_val = alpha * x_values[x_block]
-
         y_beg = y_offsets[x_col]
         y_end = y_offsets[x_col + 1]
 
-        for y_block in range(y_beg, y_end):
-            mm_block = wp.lower_bound(mm_cols, mm_beg, mm_end, y_columns[y_block])
-            mm_values[mm_block] = mm_values[mm_block] + ax_val * y_values[y_block]
+        y_block = wp.lower_bound(y_columns, y_beg, y_end, col)
+        if y_block < y_end and y_columns[y_block] == col:
+            mm_val += x_values[x_block] * y_values[y_block]
+
+    mm_values[mm_block] += alpha * mm_val
 
 
 class bsr_mm_work_arrays:
@@ -995,13 +1043,24 @@ def bsr_mm(
         kernel=_bsr_mm_count_coeffs,
         device=device,
         dim=z.nrow,
-        inputs=[copied_z_nnz, x.offsets, x.columns, y.offsets, work_arrays._mm_row_counts],
+        inputs=[
+            copied_z_nnz,
+            x.offsets,
+            x.columns,
+            y.offsets,
+            work_arrays._mm_row_counts,
+        ],
     )
     warp.utils.array_scan(work_arrays._mm_row_counts, work_arrays._mm_row_counts)
 
     # Get back total counts on host
     if device.is_cuda:
-        wp.copy(dest=work_arrays._pinned_count_buffer, src=work_arrays._mm_row_counts, src_offset=z.nrow, count=1)
+        wp.copy(
+            dest=work_arrays._pinned_count_buffer,
+            src=work_arrays._mm_row_counts,
+            src_offset=z.nrow,
+            count=1,
+        )
         wp.synchronize_stream(wp.get_stream(device))
         mm_nnz = int(work_arrays._pinned_count_buffer.numpy()[0])
     else:
@@ -1014,7 +1073,10 @@ def bsr_mm(
         # Copy z row and column indices
         wp.copy(dest=work_arrays._mm_cols, src=z.columns, count=copied_z_nnz)
         wp.launch(
-            kernel=_bsr_get_block_row, device=device, dim=copied_z_nnz, inputs=[0, z.offsets, work_arrays._mm_rows]
+            kernel=_bsr_get_block_row,
+            device=device,
+            dim=copied_z_nnz,
+            inputs=[0, z.offsets, work_arrays._mm_rows],
         )
         # Save current z values in temporary buffer
         wp.copy(src=z.values, dest=work_arrays._old_z_values, count=copied_z_nnz)
@@ -1097,7 +1159,7 @@ def bsr_mm(
     wp.launch(
         kernel=_bsr_mm_compute_values,
         device=device,
-        dim=z.nrow,
+        dim=z.nnz,
         inputs=[
             alpha,
             work_arrays._old_z_offsets if x == z else x.offsets,

warp/stubs.py

@@ -348,6 +348,78 @@ def frac(x: Float) -> Float:
     ...
 
 
+@over
+def isfinite(x: Scalar) -> bool:
+    """Return ``True`` if x is a finite number, otherwise return ``False``."""
+    ...
+
+
+@over
+def isfinite(x: Vector[Any, Scalar]) -> bool:
+    """Return ``True`` if all elements of the vector ``x`` are finite, otherwise return ``False``."""
+    ...
+
+
+@over
+def isfinite(x: Quaternion[Scalar]) -> bool:
+    """Return ``True`` if all elements of the quaternion ``x`` are finite, otherwise return ``False``."""
+    ...
+
+
+@over
+def isfinite(m: Matrix[Any, Any, Scalar]) -> bool:
+    """Return ``True`` if all elements of the matrix ``m`` are finite, otherwise return ``False``."""
+    ...
+
+
+@over
+def isnan(x: Scalar) -> bool:
+    """Return ``True`` if ``x`` is NaN, otherwise return ``False``."""
+    ...
+
+
+@over
+def isnan(x: Vector[Any, Scalar]) -> bool:
+    """Return ``True`` if any element of the vector ``x`` is NaN, otherwise return ``False``."""
+    ...
+
+
+@over
+def isnan(x: Quaternion[Scalar]) -> bool:
+    """Return ``True`` if any element of the quaternion ``x`` is NaN, otherwise return ``False``."""
+    ...
+
+
+@over
+def isnan(m: Matrix[Any, Any, Scalar]) -> bool:
+    """Return ``True`` if any element of the matrix ``m`` is NaN, otherwise return ``False``."""
+    ...
+
+
+@over
+def isinf(x: Scalar) -> bool:
+    """Return ``True`` if x is positive or negative infinity, otherwise return ``False``."""
+    ...
+
+
+@over
+def isinf(x: Vector[Any, Scalar]) -> bool:
+    """Return ``True`` if any element of the vector ``x`` is positive or negative infinity, otherwise return ``False``."""
+    ...
+
+
+@over
+def isinf(x: Quaternion[Scalar]) -> bool:
+    """Return ``True`` if any element of the quaternion ``x`` is positive or negative infinity, otherwise return ``False``."""
+    ...
+
+
+@over
+def isinf(m: Matrix[Any, Any, Scalar]) -> bool:
+    """Return ``True`` if any element of the matrix ``m`` is positive or negative infinity, otherwise return ``False``."""
+    ...
+
+
 @over
 def dot(x: Vector[Any, Scalar], y: Vector[Any, Scalar]) -> Scalar:
     """Compute the dot product between two vectors."""
@@ -1040,6 +1112,17 @@ def volume_store_i(id: uint64, i: int32, j: int32, k: int32, value: int32):
     ...
 
 
+@over
+def volume_lookup_index(id: uint64, i: int32, j: int32, k: int32) -> int32:
+    """Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
+
+    If the voxel at this index does not exist, this function returns -1.
+    This function is available for both index grids and classical volumes.
+
+    """
+    ...
+
+
 @over
 def volume_index_to_world(id: uint64, uvw: vec3f) -> vec3f:
     """Transform a point ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation."""

warp/tests/aux_test_dependent.py

@@ -10,8 +10,6 @@
 import warp as wp
 import warp.tests.aux_test_reference as ref
 
-wp.init()
-
 
 @wp.kernel
 def kern(expect: float):

warp/tests/aux_test_grad_customs.py

@@ -9,8 +9,6 @@
 
 import warp as wp
 
-wp.init()
-
 
 @wp.func
 def aux_custom_fn(x: float, y: float):

warp/tests/aux_test_reference.py

@@ -3,8 +3,6 @@
 import warp as wp
 import warp.tests.aux_test_reference_reference as refref
 
-wp.init()
-
 
 @wp.func
 def magic():

warp/tests/aux_test_reference_reference.py

@@ -2,8 +2,6 @@
 
 import warp as wp
 
-wp.init()
-
 
 @wp.func
 def more_magic():

warp/tests/aux_test_square.py

@@ -1,7 +1,5 @@
 import warp as wp
 
-wp.init()
-
 
 @wp.func
 def multiply(x: float):

warp/tests/disabled_kinematics.py

@@ -12,8 +12,6 @@ import warp as wp
 import warp.sim
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 def build_ant(num_envs):
     builder = wp.sim.ModelBuilder()

warp/tests/test_adam.py

@@ -14,8 +14,6 @@ import warp.optim
 import warp.sim
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def objective(params: wp.array(dtype=float), score: wp.array(dtype=float)):

warp/tests/test_arithmetic.py

@@ -5,7 +5,6 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import math
 import unittest
 
 import numpy as np
@@ -13,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 np_signed_int_types = [
     np.int8,
     np.int16,
@@ -838,36 +835,6 @@ def test_float_to_int(test, device, dtype, register_kernels=False):
             tape.zero()
 
 
-def test_infinity(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    def check_infinity(
-        outputs: wp.array(dtype=wptype),
-    ):
-        outputs[0] = wptype(wp.inf)
-        outputs[1] = wptype(-wp.inf)
-        outputs[2] = wptype(2.0 * wp.inf)
-        outputs[3] = wptype(-2.0 * wp.inf)
-        outputs[4] = wptype(2.0 / 0.0)
-        outputs[5] = wptype(-2.0 / 0.0)
-
-    kernel = getkernel(check_infinity, suffix=dtype.__name__)
-
-    if register_kernels:
-        return
-
-    outputs = wp.zeros(6, dtype=wptype, device=device)
-
-    wp.launch(kernel, dim=1, inputs=[], outputs=[outputs], device=device)
-
-    test.assertEqual(outputs.numpy()[0], math.inf)
-    test.assertEqual(outputs.numpy()[1], -math.inf)
-    test.assertEqual(outputs.numpy()[2], math.inf)
-    test.assertEqual(outputs.numpy()[3], -math.inf)
-    test.assertEqual(outputs.numpy()[4], math.inf)
-    test.assertEqual(outputs.numpy()[5], -math.inf)
-
-
 def test_interp(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
 
@@ -1102,9 +1069,6 @@ for dtype in np_float_types:
     add_function_test_register_kernel(
         TestArithmetic, f"test_float_to_int_{dtype.__name__}", test_float_to_int, devices=devices, dtype=dtype
     )
-    add_function_test_register_kernel(
-        TestArithmetic, f"test_infinity_{dtype.__name__}", test_infinity, devices=devices, dtype=dtype
-    )
 
 for dtype in np_scalar_types:
     add_function_test_register_kernel(

warp/tests/test_array.py

@@ -12,8 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 @wp.kernel
 def kernel_1d(a: wp.array(dtype=int, ndim=1)):
@@ -422,10 +420,10 @@ def test_view(test, device):
     wp_arr_e = wp.array(np_arr_e, dtype=wp.vec4, device=device)
     wp_arr_f = wp.array(np_arr_e, dtype=wp.quat, device=device)
 
-    assert np.array_equal(np_arr_a.view(dtype=np.float32), wp_arr_a.view(dtype=wp.float32).numpy())
-    assert np.array_equal(np_arr_b.view(dtype=np.uint32), wp_arr_b.view(dtype=wp.uint32).numpy())
-    assert np.array_equal(np_arr_c.view(dtype=np.float16), wp_arr_c.view(dtype=wp.float16).numpy())
-    assert np.array_equal(np_arr_d.view(dtype=np.uint16), wp_arr_d.view(dtype=wp.uint16).numpy())
+    assert_np_equal(wp_arr_a.view(dtype=wp.float32).numpy(), np_arr_a.view(dtype=np.float32))
+    assert_np_equal(wp_arr_b.view(dtype=wp.uint32).numpy(), np_arr_b.view(dtype=np.uint32))
+    assert_np_equal(wp_arr_c.view(dtype=wp.float16).numpy(), np_arr_c.view(dtype=np.float16))
+    assert_np_equal(wp_arr_d.view(dtype=wp.uint16).numpy(), np_arr_d.view(dtype=np.uint16))
     assert_array_equal(wp_arr_e.view(dtype=wp.quat), wp_arr_f)
 
 
@@ -486,7 +484,7 @@ def test_transpose(test, device):
     check = wp.zeros(shape=(2, 3), dtype=int, device=device)
 
     wp.launch(compare_2darrays, dim=(2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
-    assert np.array_equal(check.numpy(), np.ones((2, 3), dtype=int))
+    assert_np_equal(check.numpy(), np.ones((2, 3), dtype=int))
 
     # test transpose in square 3d case
     # wp does not support copying from/to non-contiguous arrays so check in kernel
@@ -497,7 +495,7 @@ def test_transpose(test, device):
     check = wp.zeros(shape=(3, 2, 2), dtype=int, device=device)
 
     wp.launch(compare_3darrays, dim=(3, 2, 2), inputs=[arr_transpose, arr_compare, check], device=device)
-    assert np.array_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
+    assert_np_equal(check.numpy(), np.ones((3, 2, 2), dtype=int))
 
     # test transpose in square 3d case without axes supplied
     arr_transpose = arr.transpose()
@@ -505,13 +503,13 @@ def test_transpose(test, device):
     check = wp.zeros(shape=(2, 2, 3), dtype=int, device=device)
 
     wp.launch(compare_3darrays, dim=(2, 2, 3), inputs=[arr_transpose, arr_compare, check], device=device)
-    assert np.array_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
+    assert_np_equal(check.numpy(), np.ones((2, 2, 3), dtype=int))
 
     # test transpose in 1d case (should be noop)
     np_arr = np.array([1, 2, 3], dtype=float)
     arr = wp.array(np_arr, dtype=float, device=device)
 
-    assert np.array_equal(np_arr.transpose(), arr.transpose().numpy())
+    assert_np_equal(arr.transpose().numpy(), np_arr.transpose())
 
 
 def test_fill_scalar(test, device):
@@ -2335,7 +2333,7 @@ devices = get_test_devices()
 
 class TestArray(unittest.TestCase):
     def test_array_new_del(self):
-        # test the scenario in which an array instance is created but not initalized before gc
+        # test the scenario in which an array instance is created but not initialized before gc
         instance = wp.array.__new__(wp.array)
         instance.__del__()
 

warp/tests/test_array_reduce.py

@@ -13,8 +13,6 @@ import warp as wp
 from warp.tests.unittest_utils import *
 from warp.utils import array_inner, array_sum
 
-wp.init()
-
 
 def make_test_array_sum(dtype):
     N = 1000