warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_arithmetic.py

@@ -5,9 +5,13 @@
 # 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
+
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -34,22 +38,21 @@ np_float_types = [np.float16, np.float32, np.float64]
 np_scalar_types = np_int_types + np_float_types
 
 
-def randvals(shape, dtype):
+def randvals(rng, shape, dtype):
     if dtype in np_float_types:
-        return np.random.randn(*shape).astype(dtype)
+        return rng.standard_normal(size=shape).astype(dtype)
     elif dtype in [np.int8, np.uint8, np.byte, np.ubyte]:
-        return np.random.randint(1, 3, size=shape, dtype=dtype)
-    return np.random.randint(1, 5, size=shape, dtype=dtype)
+        return rng.integers(1, high=3, size=shape, dtype=dtype)
+    return rng.integers(1, high=5, size=shape, dtype=dtype)
 
 
 kernel_cache = dict()
 
 
 def getkernel(func, suffix=""):
-    module = wp.get_module(func.__module__)
     key = func.__name__ + "_" + suffix
     if key not in kernel_cache:
-        kernel_cache[key] = wp.Kernel(func=func, key=key, module=module)
+        kernel_cache[key] = wp.Kernel(func=func, key=key)
     return kernel_cache[key]
 
 
@@ -77,7 +80,7 @@ def get_select_kernel2(dtype):
 
 
 def test_arrays(test, device, dtype):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-3,
@@ -86,14 +89,14 @@ def test_arrays(test, device, dtype):
     }.get(dtype, 0)
 
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    arr_np = randvals((10, 5), dtype)
+    arr_np = randvals(rng, (10, 5), dtype)
     arr = wp.array(arr_np, dtype=wptype, requires_grad=True, device=device)
 
     assert_np_equal(arr.numpy(), arr_np, tol=tol)
 
 
 def test_unary_ops(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -128,10 +131,12 @@ def test_unary_ops(test, device, dtype, register_kernels=False):
         return
 
     if dtype in np_float_types:
-        inputs = wp.array(np.random.randn(5, 10).astype(dtype), dtype=wptype, requires_grad=True, device=device)
+        inputs = wp.array(
+            rng.standard_normal(size=(5, 10)).astype(dtype), dtype=wptype, requires_grad=True, device=device
+        )
     else:
         inputs = wp.array(
-            np.random.randint(-2, 3, size=(5, 10), dtype=dtype), dtype=wptype, requires_grad=True, device=device
+            rng.integers(-2, high=3, size=(5, 10), dtype=dtype), dtype=wptype, requires_grad=True, device=device
         )
     outputs = wp.zeros_like(inputs)
 
@@ -207,7 +212,7 @@ def test_unary_ops(test, device, dtype, register_kernels=False):
 
 
 def test_nonzero(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -231,7 +236,7 @@ def test_nonzero(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    inputs = wp.array(np.random.randint(-2, 3, size=10).astype(dtype), dtype=wptype, requires_grad=True, device=device)
+    inputs = wp.array(rng.integers(-2, high=3, size=10).astype(dtype), dtype=wptype, requires_grad=True, device=device)
     outputs = wp.zeros_like(inputs)
 
     wp.launch(kernel, dim=1, inputs=[inputs], outputs=[outputs], device=device)
@@ -253,10 +258,10 @@ def test_nonzero(test, device, dtype, register_kernels=False):
 
 
 def test_binary_ops(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
-        np.float16: 1.0e-2,
+        np.float16: 5.0e-2,
         np.float32: 1.0e-6,
         np.float64: 1.0e-8,
     }.get(dtype, 0)
@@ -302,11 +307,11 @@ def test_binary_ops(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    vals1 = randvals([8, 10], dtype)
+    vals1 = randvals(rng, [8, 10], dtype)
     if dtype in [np_unsigned_int_types]:
-        vals2 = vals1 + randvals([8, 10], dtype)
+        vals2 = vals1 + randvals(rng, [8, 10], dtype)
     else:
-        vals2 = np.abs(randvals([8, 10], dtype))
+        vals2 = np.abs(randvals(rng, [8, 10], dtype))
 
     in1 = wp.array(vals1, dtype=wptype, requires_grad=True, device=device)
     in2 = wp.array(vals2, dtype=wptype, requires_grad=True, device=device)
@@ -458,7 +463,7 @@ def test_binary_ops(test, device, dtype, register_kernels=False):
 
 
 def test_special_funcs(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-2,
@@ -488,6 +493,7 @@ def test_special_funcs(test, device, dtype, register_kernels=False):
             outputs[11, i] = wptype(2) * wp.tanh(inputs[11, i])
             outputs[12, i] = wptype(2) * wp.acos(inputs[12, i])
             outputs[13, i] = wptype(2) * wp.asin(inputs[13, i])
+            outputs[14, i] = wptype(2) * wp.cbrt(inputs[14, i])
 
     kernel = getkernel(check_special_funcs, suffix=dtype.__name__)
     output_select_kernel = get_select_kernel2(wptype)
@@ -495,8 +501,8 @@ def test_special_funcs(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    invals = np.random.randn(14, 10).astype(dtype)
-    invals[[0, 1, 2, 7]] = 0.1 + np.abs(invals[[0, 1, 2, 7]])
+    invals = rng.normal(size=(15, 10)).astype(dtype)
+    invals[[0, 1, 2, 7, 14]] = 0.1 + np.abs(invals[[0, 1, 2, 7, 14]])
     invals[12] = np.clip(invals[12], -0.9, 0.9)
     invals[13] = np.clip(invals[13], -0.9, 0.9)
     inputs = wp.array(invals, dtype=wptype, requires_grad=True, device=device)
@@ -518,6 +524,7 @@ def test_special_funcs(test, device, dtype, register_kernels=False):
     assert_np_equal(outputs.numpy()[11], 2 * np.tanh(inputs.numpy()[11]), tol=tol)
     assert_np_equal(outputs.numpy()[12], 2 * np.arccos(inputs.numpy()[12]), tol=tol)
     assert_np_equal(outputs.numpy()[13], 2 * np.arcsin(inputs.numpy()[13]), tol=tol)
+    assert_np_equal(outputs.numpy()[14], 2 * np.cbrt(inputs.numpy()[14]), tol=tol)
 
     out = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
     if dtype in np_float_types:
@@ -694,9 +701,22 @@ def test_special_funcs(test, device, dtype, register_kernels=False):
             assert_np_equal(tape.gradients[inputs].numpy(), expected, tol=6 * tol)
             tape.zero()
 
+            # cbrt:
+            tape = wp.Tape()
+            with tape:
+                wp.launch(kernel, dim=1, inputs=[inputs], outputs=[outputs], device=device)
+                wp.launch(output_select_kernel, dim=1, inputs=[outputs, 14, i], outputs=[out], device=device)
+
+            tape.backward(loss=out)
+            expected = np.zeros_like(inputs.numpy())
+            cbrt = np.cbrt(inputs.numpy()[14, i], dtype=np.dtype(dtype))
+            expected[14, i] = (2.0 / 3.0) * (1.0 / (cbrt * cbrt))
+            assert_np_equal(tape.gradients[inputs].numpy(), expected, tol=tol)
+            tape.zero()
+
 
 def test_special_funcs_2arg(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-2,
@@ -722,8 +742,8 @@ def test_special_funcs_2arg(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    in1 = wp.array(np.abs(randvals([2, 10], dtype)), dtype=wptype, requires_grad=True, device=device)
-    in2 = wp.array(randvals([2, 10], dtype), dtype=wptype, requires_grad=True, device=device)
+    in1 = wp.array(np.abs(randvals(rng, [2, 10], dtype)), dtype=wptype, requires_grad=True, device=device)
+    in2 = wp.array(randvals(rng, [2, 10], dtype), dtype=wptype, requires_grad=True, device=device)
     outputs = wp.zeros_like(in1)
 
     wp.launch(kernel, dim=1, inputs=[in1, in2], outputs=[outputs], device=device)
@@ -763,7 +783,7 @@ def test_special_funcs_2arg(test, device, dtype, register_kernels=False):
 
 
 def test_float_to_int(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -783,6 +803,7 @@ def test_float_to_int(test, device, dtype, register_kernels=False):
             outputs[2, i] = wp.trunc(inputs[2, i])
             outputs[3, i] = wp.floor(inputs[3, i])
             outputs[4, i] = wp.ceil(inputs[4, i])
+            outputs[5, i] = wp.frac(inputs[5, i])
 
     kernel = getkernel(check_float_to_int, suffix=dtype.__name__)
     output_select_kernel = get_select_kernel2(wptype)
@@ -790,7 +811,7 @@ def test_float_to_int(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    inputs = wp.array(np.random.randn(5, 10).astype(dtype), dtype=wptype, requires_grad=True, device=device)
+    inputs = wp.array(rng.standard_normal(size=(6, 10)).astype(dtype), dtype=wptype, requires_grad=True, device=device)
     outputs = wp.zeros_like(inputs)
 
     wp.launch(kernel, dim=1, inputs=[inputs], outputs=[outputs], device=device)
@@ -800,6 +821,7 @@ def test_float_to_int(test, device, dtype, register_kernels=False):
     assert_np_equal(outputs.numpy()[2], np.trunc(inputs.numpy()[2]))
     assert_np_equal(outputs.numpy()[3], np.floor(inputs.numpy()[3]))
     assert_np_equal(outputs.numpy()[4], np.ceil(inputs.numpy()[4]))
+    assert_np_equal(outputs.numpy()[5], np.modf(inputs.numpy()[5])[0])
 
     # all the gradients should be zero as these functions are piecewise constant:
 
@@ -816,8 +838,38 @@ 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):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 1.0e-2,
@@ -844,11 +896,11 @@ def test_interp(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    e0 = randvals([2, 10], dtype)
-    e1 = e0 + randvals([2, 10], dtype) + 0.1
+    e0 = randvals(rng, [2, 10], dtype)
+    e1 = e0 + randvals(rng, [2, 10], dtype) + 0.1
     in1 = wp.array(e0, dtype=wptype, requires_grad=True, device=device)
     in2 = wp.array(e1, dtype=wptype, requires_grad=True, device=device)
-    in3 = wp.array(randvals([2, 10], dtype), dtype=wptype, requires_grad=True, device=device)
+    in3 = wp.array(randvals(rng, [2, 10], dtype), dtype=wptype, requires_grad=True, device=device)
 
     outputs = wp.zeros_like(in1)
 
@@ -948,7 +1000,7 @@ def test_interp(test, device, dtype, register_kernels=False):
 
 
 def test_clamp(test, device, dtype, register_kernels=False):
-    np.random.seed(123)
+    rng = np.random.default_rng(123)
 
     tol = {
         np.float16: 5.0e-3,
@@ -974,9 +1026,9 @@ def test_clamp(test, device, dtype, register_kernels=False):
     if register_kernels:
         return
 
-    in1 = wp.array(randvals([100], dtype), dtype=wptype, requires_grad=True, device=device)
-    starts = randvals([100], dtype)
-    diffs = np.abs(randvals([100], dtype))
+    in1 = wp.array(randvals(rng, [100], dtype), dtype=wptype, requires_grad=True, device=device)
+    starts = randvals(rng, [100], dtype)
+    diffs = np.abs(randvals(rng, [100], dtype))
     in2 = wp.array(starts, dtype=wptype, requires_grad=True, device=device)
     in3 = wp.array(starts + diffs, dtype=wptype, requires_grad=True, device=device)
     outputs = wp.zeros_like(in1)
@@ -1020,51 +1072,53 @@ def test_clamp(test, device, dtype, register_kernels=False):
             tape.zero()
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
 
-    class TestArithmetic(parent):
-        pass
 
-    # these unary ops only make sense for signed values:
-    for dtype in np_signed_int_types + np_float_types:
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_unary_ops_{dtype.__name__}", test_unary_ops, devices=devices, dtype=dtype
-        )
+class TestArithmetic(unittest.TestCase):
+    pass
 
-    for dtype in np_float_types:
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_special_funcs_{dtype.__name__}", test_special_funcs, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestArithmetic,
-            f"test_special_funcs_2arg_{dtype.__name__}",
-            test_special_funcs_2arg,
-            devices=devices,
-            dtype=dtype,
-        )
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_interp_{dtype.__name__}", test_interp, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_float_to_int_{dtype.__name__}", test_float_to_int, devices=devices, dtype=dtype
-        )
 
-    for dtype in np_scalar_types:
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_clamp_{dtype.__name__}", test_clamp, devices=devices, dtype=dtype
-        )
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_nonzero_{dtype.__name__}", test_nonzero, devices=devices, dtype=dtype
-        )
-        add_function_test(TestArithmetic, f"test_arrays_{dtype.__name__}", test_arrays, devices=devices, dtype=dtype)
-        add_function_test_register_kernel(
-            TestArithmetic, f"test_binary_ops_{dtype.__name__}", test_binary_ops, devices=devices, dtype=dtype
-        )
+# these unary ops only make sense for signed values:
+for dtype in np_signed_int_types + np_float_types:
+    add_function_test_register_kernel(
+        TestArithmetic, f"test_unary_ops_{dtype.__name__}", test_unary_ops, devices=devices, dtype=dtype
+    )
 
-    return TestArithmetic
+for dtype in np_float_types:
+    add_function_test_register_kernel(
+        TestArithmetic, f"test_special_funcs_{dtype.__name__}", test_special_funcs, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestArithmetic,
+        f"test_special_funcs_2arg_{dtype.__name__}",
+        test_special_funcs_2arg,
+        devices=devices,
+        dtype=dtype,
+    )
+    add_function_test_register_kernel(
+        TestArithmetic, f"test_interp_{dtype.__name__}", test_interp, devices=devices, dtype=dtype
+    )
+    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(
+        TestArithmetic, f"test_clamp_{dtype.__name__}", test_clamp, devices=devices, dtype=dtype
+    )
+    add_function_test_register_kernel(
+        TestArithmetic, f"test_nonzero_{dtype.__name__}", test_nonzero, devices=devices, dtype=dtype
+    )
+    add_function_test(TestArithmetic, f"test_arrays_{dtype.__name__}", test_arrays, devices=devices, dtype=dtype)
+    add_function_test_register_kernel(
+        TestArithmetic, f"test_binary_ops_{dtype.__name__}", test_binary_ops, devices=devices, dtype=dtype
+    )
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2, failfast=False)