warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_rounding.py

@@ -47,41 +47,13 @@ def test_kernel(
 
 
 def test_rounding(test, device):
-    nx = np.array(
-        [
-            4.9,
-            4.5,
-            4.1,
-            3.9,
-            3.5,
-            3.1,
-            2.9,
-            2.5,
-            2.1,
-            1.9,
-            1.5,
-            1.1,
-            0.9,
-            0.5,
-            0.1,
-            -0.1,
-            -0.5,
-            -0.9,
-            -1.1,
-            -1.5,
-            -1.9,
-            -2.1,
-            -2.5,
-            -2.9,
-            -3.1,
-            -3.5,
-            -3.9,
-            -4.1,
-            -4.5,
-            -4.9,
-        ],
-        dtype=np.float32,
-    )
+    # fmt: off
+    nx = np.array([
+        4.9,  4.5,  4.1,  3.9,  3.5,  3.1,  2.9,  2.5,  2.1,  1.9,
+        1.5,  1.1,  0.9,  0.5,  0.1, -0.1, -0.5, -0.9, -1.1, -1.5,
+       -1.9, -2.1, -2.5, -2.9, -3.1, -3.5, -3.9, -4.1, -4.5, -4.9
+    ], dtype=np.float32)
+    # fmt: on
 
     x = wp.array(nx, device=device)
     N = len(x)
@@ -149,10 +121,10 @@ def test_rounding(test, device):
     assert_np_equal(tab, golden, tol=1e-6)
 
     if print_results:
-        np.set_printoptions(formatter={"float": lambda x: "{:6.1f}".format(x).replace(".0", ".")})
+        np.set_printoptions(formatter={"float": lambda x: f"{x:6.1f}".replace(".0", ".")})
 
         print("----------------------------------------------")
-        print("   %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "cast", "floor", "ceil"))
+        print(f"   {'x ':>5s} {'round':>5s} {'rint':>5s} {'trunc':>5s} {'cast':>5s} {'floor':>5s} {'ceil':>5s}")
         print(tab)
         print("----------------------------------------------")
 
@@ -166,7 +138,7 @@ def test_rounding(test, device):
         nx_frac = np.modf(nx)[0]
 
         tab = np.stack([nx, nx_round, nx_rint, nx_trunc, nx_fix, nx_floor, nx_ceil, nx_frac], axis=1)
-        print("   %5s %5s %5s %5s %5s %5s %5s" % ("x ", "round", "rint", "trunc", "fix", "floor", "ceil"))
+        print(f"   {'x ':>5s} {'round':>5s} {'rint':>5s} {'trunc':>5s} {'fix':>5s} {'floor':>5s} {'ceil':>5s}")
         print(tab)
         print("----------------------------------------------")
 

warp/tests/test_runlength_encode.py

@@ -48,7 +48,7 @@ def test_runlength_encode_error_insufficient_storage(test, device):
     run_lengths = wp.empty(123, dtype=int, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
-        r"Output array storage sizes must be at least equal to value_count$",
+        r"Output array storage sizes must be at least equal to value_count \(123\)$",
     ):
         runlength_encode(values, run_values, run_lengths)
 
@@ -57,7 +57,7 @@ def test_runlength_encode_error_insufficient_storage(test, device):
     run_lengths = wp.empty(1, dtype=int, device="cpu")
     with test.assertRaisesRegex(
         RuntimeError,
-        r"Output array storage sizes must be at least equal to value_count$",
+        r"Output array storage sizes must be at least equal to value_count \(123\)$",
     ):
         runlength_encode(values, run_values, run_lengths)
 
@@ -68,7 +68,7 @@ def test_runlength_encode_error_dtypes_mismatch(test, device):
     run_lengths = wp.empty_like(values, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
-        r"values and run_values data types do not match$",
+        r"values and run_values data types do not match \(int32 vs float32\)$",
     ):
         runlength_encode(values, run_values, run_lengths)
 
@@ -102,7 +102,7 @@ def test_runlength_encode_error_unsupported_dtype(test, device):
     run_lengths = wp.empty(123, dtype=int, device=device)
     with test.assertRaisesRegex(
         RuntimeError,
-        r"Unsupported data type$",
+        r"Unsupported data type: float32$",
     ):
         runlength_encode(values, run_values, run_lengths)
 
@@ -118,7 +118,7 @@ class TestRunlengthEncode(unittest.TestCase):
         run_lengths = wp.empty_like(values, device="cuda:0")
         with self.assertRaisesRegex(
             RuntimeError,
-            r"Array storage devices do not match$",
+            r"run_values, run_lengths and values storage devices do not match$",
         ):
             runlength_encode(values, run_values, run_lengths)
 
@@ -127,7 +127,7 @@ class TestRunlengthEncode(unittest.TestCase):
         run_lengths = wp.empty_like(values, device="cuda:0")
         with self.assertRaisesRegex(
             RuntimeError,
-            r"Array storage devices do not match$",
+            r"run_values, run_lengths and values storage devices do not match$",
         ):
             runlength_encode(values, run_values, run_lengths)
 
@@ -136,7 +136,7 @@ class TestRunlengthEncode(unittest.TestCase):
         run_lengths = wp.empty_like(values, device="cpu")
         with self.assertRaisesRegex(
             RuntimeError,
-            r"Array storage devices do not match$",
+            r"run_values, run_lengths and values storage devices do not match$",
         ):
             runlength_encode(values, run_values, run_lengths)
 

warp/tests/test_smoothstep.py

@@ -78,7 +78,7 @@ def test_smoothstep(test, device):
             a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
             b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
             t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
-            out = wp.array([0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True)
+            out = wp.array([0] * wp.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
 
             with wp.Tape() as tape:
                 wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)

warp/tests/test_sparse.py

@@ -130,6 +130,80 @@ def test_bsr_from_triplets(test, device):
     )
     test.assertEqual(bsr.nnz, 0)
 
+    # test passing indices with wrong data ty[e]
+    rows = wp.array(rows.numpy().astype(float), dtype=float, device=device)
+    cols = wp.array(cols.numpy().astype(float), dtype=float, device=device)
+    with test.assertRaisesRegex(
+        TypeError,
+        r"Rows and columns arrays must be of type int32$",
+    ):
+        bsr_set_from_triplets(bsr, rows, cols, vals)
+
+
+def test_bsr_from_triplets_prune_numerical_zeros(test, device):
+    rows = wp.array([1, 0, 2, 3], dtype=int)
+    cols = wp.array([0, 1, 2, 3], dtype=int)
+    vals = wp.zeros(len(rows), dtype=float)
+
+    A = bsr_from_triplets(
+        rows_of_blocks=12,  # Number of rows of blocks
+        cols_of_blocks=12,  # Number of columns of blocks
+        rows=rows,  # Row indices
+        columns=cols,  # Column indices
+        values=vals,  # Block values
+        prune_numerical_zeros=False,
+    )
+    assert A.nnz_sync() == 4
+
+    A = bsr_from_triplets(
+        rows_of_blocks=12,  # Number of rows of blocks
+        cols_of_blocks=12,  # Number of columns of blocks
+        rows=rows,  # Row indices
+        columns=cols,  # Column indices
+        values=vals,  # Block values
+        prune_numerical_zeros=True,
+    )
+    assert A.nnz_sync() == 0
+
+
+def test_bsr_from_triplets_gradient(test, device):
+    rng = np.random.default_rng(123)
+
+    block_shape = (3, 3)
+    nrow = 2
+    ncol = 2
+
+    n = 4
+    rows = wp.array([1, 0, 0, 1], dtype=int, device=device)
+    cols = wp.array([0, 1, 0, 0], dtype=int, device=device)
+
+    vals = wp.array(
+        rng.random(size=(n, block_shape[0], block_shape[1])), dtype=wp.mat33, device=device, requires_grad=True
+    )
+
+    with wp.Tape() as tape:
+        mat = bsr_from_triplets(nrow, ncol, rows, cols, vals)
+
+    assert mat.nnz_sync() == 3
+
+    zero_block = np.zeros((3, 3))
+    ones_block = np.ones((3, 3))
+
+    mat.values.grad[0:1].fill_(1.0)
+    tape.backward()
+    assert_np_equal(vals.grad.numpy(), np.stack((zero_block, zero_block, ones_block, zero_block)))
+    tape.zero()
+
+    mat.values.grad[1:2].fill_(1.0)
+    tape.backward()
+    assert_np_equal(vals.grad.numpy(), np.stack((zero_block, ones_block, zero_block, zero_block)))
+    tape.zero()
+
+    mat.values.grad[2:3].fill_(1.0)
+    tape.backward()
+    assert_np_equal(vals.grad.numpy(), np.stack((ones_block, zero_block, zero_block, ones_block)))
+    tape.zero()
+
 
 def test_bsr_get_set_diag(test, device):
     rng = np.random.default_rng(123)
@@ -191,7 +265,7 @@ def test_bsr_get_set_diag(test, device):
     assert_np_equal(diag_bsr.values.numpy(), np.broadcast_to(np.eye(4), shape=(nrow, 4, 4)), tol=0.000001)
 
     diag_csr = bsr_identity(nrow, block_type=wp.float64, device=device)
-    assert np.all(diag_csr.values.numpy() == np.ones(nrow, dtype=float))
+    np.testing.assert_array_equal(diag_csr.values.numpy(), np.ones(nrow, dtype=float))
 
 
 def test_bsr_split_merge(test, device):
@@ -239,7 +313,7 @@ def test_bsr_split_merge(test, device):
 
     with test.assertRaisesRegex(
         ValueError,
-        "The requested block shape does not evenly divide the source matrix",
+        r"The requested block shape \(32, 32\) does not evenly divide the source matrix of total size \(16, 16\)",
     ):
         bsr_copy(bsr, block_shape=(32, 32))
 
@@ -556,9 +630,16 @@ class TestSparse(unittest.TestCase):
 
 add_function_test(TestSparse, "test_csr_from_triplets", test_csr_from_triplets, devices=devices)
 add_function_test(TestSparse, "test_bsr_from_triplets", test_bsr_from_triplets, devices=devices)
+add_function_test(
+    TestSparse,
+    "test_bsr_from_triplets_prune_numerical_zeros",
+    test_bsr_from_triplets_prune_numerical_zeros,
+    devices=devices,
+)
 add_function_test(TestSparse, "test_bsr_get_diag", test_bsr_get_set_diag, devices=devices)
 add_function_test(TestSparse, "test_bsr_split_merge", test_bsr_split_merge, devices=devices)
 add_function_test(TestSparse, "test_bsr_assign_masked", test_bsr_assign_masked, devices=devices)
+add_function_test(TestSparse, "test_bsr_from_triplets_gradient", test_bsr_from_triplets_gradient, devices=devices)
 
 add_function_test(TestSparse, "test_csr_transpose", make_test_bsr_transpose((1, 1), wp.float32), devices=devices)
 add_function_test(TestSparse, "test_bsr_transpose_1_3", make_test_bsr_transpose((1, 3), wp.float32), devices=devices)