warp-lang 1.4.1__py3-none-manylinux2014_aarch64.whl → 1.4.2__py3-none-manylinux2014_aarch64.whl

warp/tests/test_array.py

@@ -2361,64 +2361,75 @@ def test_array_from_cai(test, device):
     assert_np_equal(arr_warp.numpy(), np.array([[2, 1, 1], [1, 0, 0], [1, 0, 0]]))
 
 
-def test_array_inplace_ops(test, device):
-    @wp.kernel
-    def inplace_add_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
-        i = wp.tid()
-        x[i] += y[i]
+@wp.kernel
+def inplace_add_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    i = wp.tid()
+    x[i] += y[i]
 
-    @wp.kernel
-    def inplace_add_2d(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
-        i, j = wp.tid()
-        x[i, j] += y[i, j]
 
-    @wp.kernel
-    def inplace_add_3d(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float)):
-        i, j, k = wp.tid()
-        x[i, j, k] += y[i, j, k]
+@wp.kernel
+def inplace_add_2d(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+    x[i, j] += y[i, j]
 
-    @wp.kernel
-    def inplace_add_4d(x: wp.array4d(dtype=float), y: wp.array4d(dtype=float)):
-        i, j, k, l = wp.tid()
-        x[i, j, k, l] += y[i, j, k, l]
 
-    @wp.kernel
-    def inplace_sub_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
-        i = wp.tid()
-        x[i] -= y[i]
+@wp.kernel
+def inplace_add_3d(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float)):
+    i, j, k = wp.tid()
+    x[i, j, k] += y[i, j, k]
 
-    @wp.kernel
-    def inplace_sub_2d(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
-        i, j = wp.tid()
-        x[i, j] -= y[i, j]
 
-    @wp.kernel
-    def inplace_sub_3d(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float)):
-        i, j, k = wp.tid()
-        x[i, j, k] -= y[i, j, k]
+@wp.kernel
+def inplace_add_4d(x: wp.array4d(dtype=float), y: wp.array4d(dtype=float)):
+    i, j, k, l = wp.tid()
+    x[i, j, k, l] += y[i, j, k, l]
 
-    @wp.kernel
-    def inplace_sub_4d(x: wp.array4d(dtype=float), y: wp.array4d(dtype=float)):
-        i, j, k, l = wp.tid()
-        x[i, j, k, l] -= y[i, j, k, l]
 
-    @wp.kernel
-    def inplace_add_vecs(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
-        i = wp.tid()
-        x[i] += y[i]
+@wp.kernel
+def inplace_sub_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    i = wp.tid()
+    x[i] -= y[i]
 
-    @wp.kernel
-    def inplace_add_mats(x: wp.array(dtype=wp.mat33), y: wp.array(dtype=wp.mat33)):
-        i = wp.tid()
-        x[i] += y[i]
 
-    @wp.kernel
-    def inplace_add_rhs(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
-        i = wp.tid()
-        a = y[i]
-        a += x[i]
-        wp.atomic_add(z, 0, a)
+@wp.kernel
+def inplace_sub_2d(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+    x[i, j] -= y[i, j]
+
+
+@wp.kernel
+def inplace_sub_3d(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float)):
+    i, j, k = wp.tid()
+    x[i, j, k] -= y[i, j, k]
+
+
+@wp.kernel
+def inplace_sub_4d(x: wp.array4d(dtype=float), y: wp.array4d(dtype=float)):
+    i, j, k, l = wp.tid()
+    x[i, j, k, l] -= y[i, j, k, l]
+
 
+@wp.kernel
+def inplace_add_vecs(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=wp.vec3)):
+    i = wp.tid()
+    x[i] += y[i]
+
+
+@wp.kernel
+def inplace_add_mats(x: wp.array(dtype=wp.mat33), y: wp.array(dtype=wp.mat33)):
+    i = wp.tid()
+    x[i] += y[i]
+
+
+@wp.kernel
+def inplace_add_rhs(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+    i = wp.tid()
+    a = y[i]
+    a += x[i]
+    wp.atomic_add(z, 0, a)
+
+
+def test_array_inplace_diff_ops(test, device):
     N = 3
     x1 = wp.ones(N, dtype=float, requires_grad=True, device=device)
     x2 = wp.ones((N, N), dtype=float, requires_grad=True, device=device)
@@ -2528,6 +2539,32 @@ def test_array_inplace_ops(test, device):
     assert_np_equal(y.grad.numpy(), np.ones(1, dtype=float))
 
 
+@wp.kernel
+def inplace_mul_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    i = wp.tid()
+    x[i] *= y[i]
+
+
+@wp.kernel
+def inplace_div_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    i = wp.tid()
+    x[i] /= y[i]
+
+
+def test_array_inplace_non_diff_ops(test, device):
+    N = 3
+    x1 = wp.full(N, value=10.0, dtype=float, device=device)
+    y1 = wp.full(N, value=5.0, dtype=float, device=device)
+
+    wp.launch(inplace_mul_1d, N, inputs=[x1, y1], device=device)
+    assert_np_equal(x1.numpy(), np.full(N, fill_value=50.0, dtype=float))
+
+    x1.fill_(10.0)
+    y1.fill_(5.0)
+    wp.launch(inplace_div_1d, N, inputs=[x1, y1], device=device)
+    assert_np_equal(x1.numpy(), np.full(N, fill_value=2.0, dtype=float))
+
+
 @wp.kernel
 def inc_scalar(a: wp.array(dtype=float)):
     tid = wp.tid()
@@ -2609,6 +2646,87 @@ def test_numpy_array_interface(test, device):
         assert a1.strides == a2.strides
 
 
+@wp.kernel
+def kernel_indexing_types(
+    arr_1d: wp.array(dtype=wp.int32, ndim=1),
+    arr_2d: wp.array(dtype=wp.int32, ndim=2),
+    arr_3d: wp.array(dtype=wp.int32, ndim=3),
+    arr_4d: wp.array(dtype=wp.int32, ndim=4),
+):
+    x = arr_1d[wp.uint8(0)]
+    y = arr_1d[wp.int16(1)]
+    z = arr_1d[wp.uint32(2)]
+    w = arr_1d[wp.int64(3)]
+
+    x = arr_2d[wp.uint8(0), wp.uint8(0)]
+    y = arr_2d[wp.int16(1), wp.int16(1)]
+    z = arr_2d[wp.uint32(2), wp.uint32(2)]
+    w = arr_2d[wp.int64(3), wp.int64(3)]
+
+    x = arr_3d[wp.uint8(0), wp.uint8(0), wp.uint8(0)]
+    y = arr_3d[wp.int16(1), wp.int16(1), wp.int16(1)]
+    z = arr_3d[wp.uint32(2), wp.uint32(2), wp.uint32(2)]
+    w = arr_3d[wp.int64(3), wp.int64(3), wp.int64(3)]
+
+    x = arr_4d[wp.uint8(0), wp.uint8(0), wp.uint8(0), wp.uint8(0)]
+    y = arr_4d[wp.int16(1), wp.int16(1), wp.int16(1), wp.int16(1)]
+    z = arr_4d[wp.uint32(2), wp.uint32(2), wp.uint32(2), wp.uint32(2)]
+    w = arr_4d[wp.int64(3), wp.int64(3), wp.int64(3), wp.int64(3)]
+
+    arr_1d[wp.uint8(0)] = 123
+    arr_1d[wp.int16(1)] = 123
+    arr_1d[wp.uint32(2)] = 123
+    arr_1d[wp.int64(3)] = 123
+
+    arr_2d[wp.uint8(0), wp.uint8(0)] = 123
+    arr_2d[wp.int16(1), wp.int16(1)] = 123
+    arr_2d[wp.uint32(2), wp.uint32(2)] = 123
+    arr_2d[wp.int64(3), wp.int64(3)] = 123
+
+    arr_3d[wp.uint8(0), wp.uint8(0), wp.uint8(0)] = 123
+    arr_3d[wp.int16(1), wp.int16(1), wp.int16(1)] = 123
+    arr_3d[wp.uint32(2), wp.uint32(2), wp.uint32(2)] = 123
+    arr_3d[wp.int64(3), wp.int64(3), wp.int64(3)] = 123
+
+    arr_4d[wp.uint8(0), wp.uint8(0), wp.uint8(0), wp.uint8(0)] = 123
+    arr_4d[wp.int16(1), wp.int16(1), wp.int16(1), wp.int16(1)] = 123
+    arr_4d[wp.uint32(2), wp.uint32(2), wp.uint32(2), wp.uint32(2)] = 123
+    arr_4d[wp.int64(3), wp.int64(3), wp.int64(3), wp.int64(3)] = 123
+
+    wp.atomic_add(arr_1d, wp.uint8(0), 123)
+    wp.atomic_sub(arr_1d, wp.int16(1), 123)
+    wp.atomic_min(arr_1d, wp.uint32(2), 123)
+    wp.atomic_max(arr_1d, wp.int64(3), 123)
+
+    wp.atomic_add(arr_2d, wp.uint8(0), wp.uint8(0), 123)
+    wp.atomic_sub(arr_2d, wp.int16(1), wp.int16(1), 123)
+    wp.atomic_min(arr_2d, wp.uint32(2), wp.uint32(2), 123)
+    wp.atomic_max(arr_2d, wp.int64(3), wp.int64(3), 123)
+
+    wp.atomic_add(arr_3d, wp.uint8(0), wp.uint8(0), wp.uint8(0), 123)
+    wp.atomic_sub(arr_3d, wp.int16(1), wp.int16(1), wp.int16(1), 123)
+    wp.atomic_min(arr_3d, wp.uint32(2), wp.uint32(2), wp.uint32(2), 123)
+    wp.atomic_max(arr_3d, wp.int64(3), wp.int64(3), wp.int64(3), 123)
+
+    wp.atomic_add(arr_4d, wp.uint8(0), wp.uint8(0), wp.uint8(0), wp.uint8(0), 123)
+    wp.atomic_sub(arr_4d, wp.int16(1), wp.int16(1), wp.int16(1), wp.int16(1), 123)
+    wp.atomic_min(arr_4d, wp.uint32(2), wp.uint32(2), wp.uint32(2), wp.uint32(2), 123)
+    wp.atomic_max(arr_4d, wp.int64(3), wp.int64(3), wp.int64(3), wp.int64(3), 123)
+
+
+def test_indexing_types(test, device):
+    arr_1d = wp.zeros(shape=(4,), dtype=wp.int32, device=device)
+    arr_2d = wp.zeros(shape=(4, 4), dtype=wp.int32, device=device)
+    arr_3d = wp.zeros(shape=(4, 4, 4), dtype=wp.int32, device=device)
+    arr_4d = wp.zeros(shape=(4, 4, 4, 4), dtype=wp.int32, device=device)
+    wp.launch(
+        kernel=kernel_indexing_types,
+        dim=1,
+        inputs=(arr_1d, arr_2d, arr_3d, arr_4d),
+        device=device,
+    )
+
+
 devices = get_test_devices()
 
 
@@ -2669,12 +2787,14 @@ add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, dev
 add_function_test(TestArray, "test_array_aliasing_from_numpy", test_array_aliasing_from_numpy, devices=["cpu"])
 add_function_test(TestArray, "test_numpy_array_interface", test_numpy_array_interface, devices=["cpu"])
 
-add_function_test(TestArray, "test_array_inplace_ops", test_array_inplace_ops, devices=devices)
+add_function_test(TestArray, "test_array_inplace_diff_ops", test_array_inplace_diff_ops, devices=devices)
+add_function_test(TestArray, "test_array_inplace_non_diff_ops", test_array_inplace_non_diff_ops, devices=devices)
 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)
+add_function_test(TestArray, "test_indexing_types", test_indexing_types, devices=devices)
 
 try:
     import torch

warp/tests/test_closest_point_edge_edge.py

@@ -220,12 +220,12 @@ def check_edge_closest_point_sufficient_necessary_kernel(
 
 def check_edge_closest_point_random(test, device):
     num_tests = 100000
-    np.random.seed(12345)
-    p1 = wp.array(np.random.randn(num_tests, 3), dtype=wp.vec3, device=device)
-    q1 = wp.array(np.random.randn(num_tests, 3), dtype=wp.vec3, device=device)
+    rng = np.random.default_rng(123)
+    p1 = wp.array(rng.standard_normal(size=(num_tests, 3)), dtype=wp.vec3, device=device)
+    q1 = wp.array(rng.standard_normal(size=(num_tests, 3)), dtype=wp.vec3, device=device)
 
-    p2 = wp.array(np.random.randn(num_tests, 3), dtype=wp.vec3, device=device)
-    q2 = wp.array(np.random.randn(num_tests, 3), dtype=wp.vec3, device=device)
+    p2 = wp.array(rng.standard_normal(size=(num_tests, 3)), dtype=wp.vec3, device=device)
+    q2 = wp.array(rng.standard_normal(size=(num_tests, 3)), dtype=wp.vec3, device=device)
 
     wp.launch(
         kernel=check_edge_closest_point_sufficient_necessary_kernel,
@@ -235,10 +235,10 @@ def check_edge_closest_point_random(test, device):
     )
 
     # parallel edges
-    p1 = np.random.randn(num_tests, 3)
-    q1 = np.random.randn(num_tests, 3)
+    p1 = rng.standard_normal(size=(num_tests, 3))
+    q1 = rng.standard_normal(size=(num_tests, 3))
 
-    shifts = np.random.randn(num_tests, 3)
+    shifts = rng.standard_normal(size=(num_tests, 3))
 
     p2 = p1 + shifts
     q2 = q1 + shifts

warp/tests/test_codegen.py

@@ -534,6 +534,45 @@ def test_error_mutating_constant_in_dynamic_loop(test, device):
     )
     assert_np_equal(output.numpy(), np.ones([num_threads, const_a + const_b + dyn_a + dyn_b + dyn_c + 1]))
 
+    @wp.kernel
+    def static_then_dynamic_loop_kernel(mats: wp.array(dtype=wp.mat33d)):
+        tid = wp.tid()
+        mat = wp.mat33d()
+        for i in range(3):
+            for j in range(3):
+                mat[i, j] = wp.float64(0.0)
+
+        dim = 2
+        for i in range(dim + 1):
+            for j in range(dim + 1):
+                mat[i, j] = wp.float64(1.0)
+
+        mats[tid] = mat
+
+    mats = wp.empty(1, dtype=wp.mat33d, device=device)
+    wp.launch(static_then_dynamic_loop_kernel, dim=1, inputs=[mats], device=device)
+    assert_np_equal(mats.numpy(), np.ones((1, 3, 3)))
+
+    @wp.kernel
+    def dynamic_then_static_loop_kernel(mats: wp.array(dtype=wp.mat33d)):
+        tid = wp.tid()
+        mat = wp.mat33d()
+
+        dim = 2
+        for i in range(dim + 1):
+            for j in range(dim + 1):
+                mat[i, j] = wp.float64(1.0)
+
+        for i in range(3):
+            for j in range(3):
+                mat[i, j] = wp.float64(0.0)
+
+        mats[tid] = mat
+
+    mats = wp.empty(1, dtype=wp.mat33d, device=device)
+    wp.launch(dynamic_then_static_loop_kernel, dim=1, inputs=[mats], device=device)
+    assert_np_equal(mats.numpy(), np.zeros((1, 3, 3)))
+
 
 @wp.kernel
 def test_call_syntax():

warp/tests/test_fabricarray.py

@@ -821,6 +821,38 @@ def test_fabricarray_fill_matrix(test, device):
             assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
 
 
+@wp.kernel
+def fa_kernel_indexing_types(
+    a: wp.fabricarray(dtype=wp.int32),
+):
+    x = a[wp.uint8(0)]
+    y = a[wp.int16(1)]
+    z = a[wp.uint32(2)]
+    w = a[wp.int64(3)]
+
+    a[wp.uint8(0)] = 123
+    a[wp.int16(1)] = 123
+    a[wp.uint32(2)] = 123
+    a[wp.int64(3)] = 123
+
+    wp.atomic_add(a, wp.uint8(0), 123)
+    wp.atomic_sub(a, wp.int16(1), 123)
+    # wp.atomic_min(a, wp.uint32(2), 123)
+    # wp.atomic_max(a, wp.int64(3), 123)
+
+
+def test_fabricarray_indexing_types(test, device):
+    data = wp.zeros(shape=(4,), dtype=wp.int32, device=device)
+    iface = _create_fabric_array_interface(data, "foo", copy=True)
+    fa = wp.fabricarray(data=iface, attrib="foo")
+    wp.launch(
+        kernel=fa_kernel_indexing_types,
+        dim=1,
+        inputs=(fa,),
+        device=device,
+    )
+
+
 @wp.kernel
 def fa_generic_sums_kernel(a: wp.fabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
     i = wp.tid()
@@ -945,6 +977,7 @@ add_function_test(TestFabricArray, "test_fabricarray_generic_array", test_fabric
 add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
 add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
 add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)
+add_function_test(TestFabricArray, "test_fabricarray_indexing_types", test_fabricarray_indexing_types, devices=devices)
 
 # fabric arrays of arrays
 add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)

warp/tests/test_func.py

@@ -7,7 +7,7 @@
 
 import math
 import unittest
-from typing import Tuple
+from typing import Any, Tuple
 
 import numpy as np
 
@@ -191,6 +191,37 @@ def test_user_func_return_multiple_values():
     wp.expect_eq(b, 54756.0)
 
 
+@wp.func
+def user_func_overload(
+    b: wp.array(dtype=Any),
+    i: int,
+):
+    return b[i] * 2.0
+
+
+@wp.kernel
+def user_func_overload_resolution_kernel(
+    a: wp.array(dtype=Any),
+    b: wp.array(dtype=Any),
+):
+    i = wp.tid()
+    a[i] = user_func_overload(b, i)
+
+
+def test_user_func_overload_resolution(test, device):
+    a0 = wp.array((1, 2, 3), dtype=wp.vec3)
+    b0 = wp.array((2, 3, 4), dtype=wp.vec3)
+
+    a1 = wp.array((5,), dtype=float)
+    b1 = wp.array((6,), dtype=float)
+
+    wp.launch(user_func_overload_resolution_kernel, a0.shape, (a0, b0))
+    wp.launch(user_func_overload_resolution_kernel, a1.shape, (a1, b1))
+
+    assert_np_equal(a0.numpy()[0], (4, 6, 8))
+    assert a1.numpy()[0] == 12
+
+
 devices = get_test_devices()
 
 
@@ -375,6 +406,9 @@ add_kernel_test(
     dim=1,
     devices=devices,
 )
+add_function_test(
+    TestFunc, func=test_user_func_overload_resolution, name="test_user_func_overload_resolution", devices=devices
+)
 
 
 if __name__ == "__main__":

warp/tests/test_mesh_query_point.py

@@ -769,20 +769,21 @@ def point_query_aabb_and_closest(
 
 @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
 def test_set_mesh_points(test, device):
+    rng = np.random.default_rng(123)
+
     vs, fs = load_mesh()
 
     vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
-    velocities1_np = np.random.randn(vertices1.shape[0], 3)
+    velocities1_np = rng.standard_normal(size=(vertices1.shape[0], 3))
     velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
 
     faces = wp.array(fs, dtype=wp.int32, device=device)
     mesh = wp.Mesh(vertices1, faces, velocities=velocities1)
     fs_2D = faces.reshape((-1, 3))
-    np.random.seed(12345)
     n = 1000
     query_radius = 0.2
 
-    pts1 = wp.array(np.random.randn(n, 3), dtype=wp.vec3, device=device)
+    pts1 = wp.array(rng.standard_normal(size=(n, 3)), dtype=wp.vec3, device=device)
 
     query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
     query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)

warp/tests/test_print.py

@@ -7,6 +7,7 @@
 
 import sys
 import unittest
+from typing import Any
 
 import warp as wp
 from warp.tests.unittest_utils import *
@@ -126,6 +127,139 @@ def test_print_boolean(test, device):
         test.assertRegex(s, rf"True{os.linesep}False{os.linesep}")
 
 
+@wp.kernel
+def generic_print_kernel(x: Any):
+    print(x)
+
+
+@wp.struct
+class SimpleStruct:
+    x: float
+    y: float
+
+
+generic_print_types = [*wp.types.scalar_types]
+for scalar_type in wp.types.scalar_types:
+    generic_print_types.append(wp.types.vector(2, scalar_type))
+    generic_print_types.append(wp.types.vector(3, scalar_type))
+    generic_print_types.append(wp.types.vector(4, scalar_type))
+    generic_print_types.append(wp.types.matrix((2, 2), scalar_type))
+    generic_print_types.append(wp.types.matrix((3, 3), scalar_type))
+    generic_print_types.append(wp.types.matrix((4, 4), scalar_type))
+generic_print_types.append(wp.bool)
+generic_print_types.append(SimpleStruct)
+generic_print_types.append(wp.array(dtype=float))
+
+for T in generic_print_types:
+    wp.overload(generic_print_kernel, [T])
+
+
+def test_print_adjoint(test, device):
+    for scalar_type in wp.types.scalar_types:
+        # scalar
+        capture = StdOutCapture()
+        capture.begin()
+        wp.launch(
+            generic_print_kernel,
+            dim=1,
+            inputs=[scalar_type(17)],
+            adj_inputs=[scalar_type(42)],
+            adjoint=True,
+            device=device,
+        )
+        wp.synchronize_device(device)
+        s = capture.end()
+
+        # We skip the win32 comparison for now since the capture sometimes is an empty string
+        if sys.platform != "win32":
+            test.assertRegex(s, rf"17{os.linesep}adj: 42{os.linesep}")
+
+        for dim in (2, 3, 4):
+            # vector
+            vec_type = wp.types.vector(dim, scalar_type)
+            vec_data = np.arange(vec_type._length_, dtype=wp.dtype_to_numpy(scalar_type))
+            v = vec_type(vec_data)
+            adj_v = vec_type(vec_data[::-1])
+
+            capture = StdOutCapture()
+            capture.begin()
+            wp.launch(generic_print_kernel, dim=1, inputs=[v], adj_inputs=[adj_v], adjoint=True, device=device)
+            wp.synchronize_device(device)
+            s = capture.end()
+
+            # We skip the win32 comparison for now since the capture sometimes is an empty string
+            if sys.platform != "win32":
+                expected_forward = " ".join(str(int(x)) for x in v) + " "
+                expected_adjoint = " ".join(str(int(x)) for x in adj_v)
+                test.assertRegex(s, rf"{expected_forward}{os.linesep}adj: {expected_adjoint}{os.linesep}")
+
+            # matrix
+            mat_type = wp.types.matrix((dim, dim), scalar_type)
+            mat_data = np.arange(mat_type._length_, dtype=wp.dtype_to_numpy(scalar_type))
+            m = mat_type(mat_data)
+            adj_m = mat_type(mat_data[::-1])
+
+            capture = StdOutCapture()
+            capture.begin()
+            wp.launch(generic_print_kernel, dim=1, inputs=[m], adj_inputs=[adj_m], adjoint=True, device=device)
+            wp.synchronize_device(device)
+            s = capture.end()
+
+            # We skip the win32 comparison for now since the capture sometimes is an empty string
+            if sys.platform != "win32":
+                expected_forward = ""
+                expected_adjoint = ""
+                for row in range(dim):
+                    if row == 0:
+                        adj_prefix = "adj: "
+                    else:
+                        adj_prefix = "     "
+                    expected_forward += " ".join(str(int(x)) for x in m[row]) + f" {os.linesep}"
+                    expected_adjoint += adj_prefix + " ".join(str(int(x)) for x in adj_m[row]) + f"{os.linesep}"
+                test.assertRegex(s, rf"{expected_forward}{expected_adjoint}")
+
+    # Booleans
+    capture = StdOutCapture()
+    capture.begin()
+    wp.launch(generic_print_kernel, dim=1, inputs=[True], adj_inputs=[False], adjoint=True, device=device)
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(s, rf"True{os.linesep}adj: False{os.linesep}")
+
+    # structs, not printable yet
+    capture = StdOutCapture()
+    capture.begin()
+    wp.launch(
+        generic_print_kernel, dim=1, inputs=[SimpleStruct()], adj_inputs=[SimpleStruct()], adjoint=True, device=device
+    )
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(
+            s, rf"<type without print implementation>{os.linesep}adj: <type without print implementation>{os.linesep}"
+        )
+
+    # arrays, not printable
+    capture = StdOutCapture()
+    capture.begin()
+    a = wp.ones(10, dtype=float, device=device)
+    adj_a = wp.zeros(10, dtype=float, device=device)
+    wp.launch(generic_print_kernel, dim=1, inputs=[a], adj_inputs=[adj_a], adjoint=True, device=device)
+    wp.synchronize_device(device)
+    s = capture.end()
+
+    # We skip the win32 comparison for now since the capture sometimes is an empty string
+    if sys.platform != "win32":
+        test.assertRegex(
+            s, rf"<type without print implementation>{os.linesep}adj: <type without print implementation>{os.linesep}"
+        )
+
+
 class TestPrint(unittest.TestCase):
     pass
 
@@ -134,6 +268,7 @@ devices = get_test_devices()
 add_function_test(TestPrint, "test_print", test_print, devices=devices, check_output=False)
 add_function_test(TestPrint, "test_print_numeric", test_print_numeric, devices=devices, check_output=False)
 add_function_test(TestPrint, "test_print_boolean", test_print_boolean, devices=devices, check_output=False)
+add_function_test(TestPrint, "test_print_adjoint", test_print_adjoint, devices=devices, check_output=False)
 
 
 if __name__ == "__main__":

warp/tests/unittest_suites.py

@@ -170,6 +170,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_sparse import TestSparse
     from warp.tests.test_spatial import TestSpatial
     from warp.tests.test_special_values import TestSpecialValues
+    from warp.tests.test_static import TestStatic
     from warp.tests.test_streams import TestStreams
     from warp.tests.test_struct import TestStruct
     from warp.tests.test_tape import TestTape
@@ -269,6 +270,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestSparse,
         TestSpatial,
         TestSpecialValues,
+        TestStatic,
         TestStreams,
         TestStruct,
         TestTape,
@@ -329,6 +331,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     from warp.tests.test_rounding import TestRounding
     from warp.tests.test_runlength_encode import TestRunlengthEncode
     from warp.tests.test_sparse import TestSparse
+    from warp.tests.test_static import TestStatic
     from warp.tests.test_streams import TestStreams
     from warp.tests.test_tape import TestTape
     from warp.tests.test_transient_module import TestTransientModule
@@ -374,6 +377,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestRounding,
         TestRunlengthEncode,
         TestSparse,
+        TestStatic,
         TestStreams,
         TestTape,
         TestTransientModule,

warp/types.py

@@ -1488,7 +1488,7 @@ def types_equal(a, b, match_generic=False):
 
         return True
 
-    if is_array(a) and type(a) is type(b):
+    if is_array(a) and type(a) is type(b) and types_equal(a.dtype, b.dtype, match_generic=match_generic):
         return True
 
     # match NewStructInstance and Struct dtype

{warp_lang-1.4.1.dist-info → warp_lang-1.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: warp-lang
-Version: 1.4.1
+Version: 1.4.2
 Summary: A Python framework for high-performance simulation and graphics programming
 Author-email: NVIDIA Corporation <mmacklin@nvidia.com>
 License: NVIDIA Software License
@@ -77,9 +77,9 @@ the `pip install` command, e.g.
 
 | Platform        | Install Command                                                                                                               |
 | --------------- | ----------------------------------------------------------------------------------------------------------------------------- |
-| Linux aarch64   | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.1/warp_lang-1.4.1+cu11-py3-none-manylinux2014_aarch64.whl` |
-| Linux x86-64    | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.1/warp_lang-1.4.1+cu11-py3-none-manylinux2014_x86_64.whl`  |
-| Windows x86-64  | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.1/warp_lang-1.4.1+cu11-py3-none-win_amd64.whl`             |
+| Linux aarch64   | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.2/warp_lang-1.4.2+cu11-py3-none-manylinux2014_aarch64.whl` |
+| Linux x86-64    | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.2/warp_lang-1.4.2+cu11-py3-none-manylinux2014_x86_64.whl`  |
+| Windows x86-64  | `pip install https://github.com/NVIDIA/warp/releases/download/v1.4.2/warp_lang-1.4.2+cu11-py3-none-win_amd64.whl`             |
 
 The `--force-reinstall` option may need to be used to overwrite a previous installation.