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

warp/tape.py

@@ -5,11 +5,39 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
 import warp as wp
 
 
 class Tape:
+    """
+    Record kernel launches within a Tape scope to enable automatic differentiation.
+    Gradients can be computed after the operations have been recorded on the tape via
+    ``tape.backward()``.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        tape = wp.Tape()
+
+        # forward pass
+        with tape:
+            wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
+            wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
+            wp.launch(kernel=loss, inputs=[d, l], device="cuda")
+
+        # reverse pass
+        tape.backward(l)
+
+    Gradients can be accessed via the ``tape.gradients`` dictionary, e.g.:
+
+    .. code-block:: python
+
+        print(tape.gradients[a])
+
+    """
+
     def __init__(self):
         self.gradients = {}
         self.const_gradients = set()
@@ -37,13 +65,24 @@ class Tape:
     #  adj_tensor = tape.gradients[tensor]
     #
     def backward(self, loss: wp.array = None, grads: dict = None):
+        """
+        Evaluate the backward pass of the recorded operations on the tape.
+        A single-element array ``loss`` or a dictionary of arrays ``grads``
+        can be provided to assign the incoming gradients for the reverse-mode
+        automatic differentiation pass.
+
+        Args:
+            loss (wp.array): A single-element array that holds the loss function value whose gradient is to be computed
+            grads (dict): A dictionary of arrays that map from Warp arrays to their incoming gradients
+
+        """
         # if scalar loss is specified then initialize
         # a 'seed' array for it, with gradient of one
         if loss:
             if loss.size > 1 or wp.types.type_length(loss.dtype) > 1:
                 raise RuntimeError("Can only return gradients for scalar loss functions.")
 
-            if loss.requires_grad == False:
+            if not loss.requires_grad:
                 raise RuntimeError(
                     "Scalar loss arrays should have requires_grad=True set before calling Tape.backward()"
                 )
@@ -67,9 +106,10 @@ class Tape:
             else:
                 kernel = launch[0]
                 dim = launch[1]
-                inputs = launch[2]
-                outputs = launch[3]
-                device = launch[4]
+                max_blocks = launch[2]
+                inputs = launch[3]
+                outputs = launch[4]
+                device = launch[5]
 
                 adj_inputs = []
                 adj_outputs = []
@@ -91,16 +131,21 @@ class Tape:
                     adj_outputs=adj_outputs,
                     device=device,
                     adjoint=True,
+                    max_blocks=max_blocks,
                 )
 
     # record a kernel launch on the tape
-    def record_launch(self, kernel, dim, inputs, outputs, device):
-        self.launches.append([kernel, dim, inputs, outputs, device])
+    def record_launch(self, kernel, dim, max_blocks, inputs, outputs, device):
+        self.launches.append([kernel, dim, max_blocks, inputs, outputs, device])
 
-    # records a custom function for the backward pass, can be any
-    # Callable python object. Callee should also pass arrays that
-    # take part in the function for gradient tracking.
     def record_func(self, backward, arrays):
+        """
+        Records a custom function to be executed only in the backward pass.
+
+        Args:
+            backward (Callable): A callable Python object (can be any function) that will be executed in the backward pass.
+            arrays (list): A list of arrays that are used by the function for gradient tracking.
+        """
         self.launches.append(backward)
 
         for a in arrays:
@@ -125,11 +170,11 @@ class Tape:
             return a.grad
 
         elif isinstance(a, wp.codegen.StructInstance):
-            adj = a._struct_()
-            for name, _ in a._struct_.ctype._fields_:
+            adj = a._cls()
+            for name, _ in a._cls.ctype._fields_:
                 if name.startswith("_"):
                     continue
-                if isinstance(a._struct_.vars[name].type, wp.array):
+                if isinstance(a._cls.vars[name].type, wp.array):
                     arr = getattr(a, name)
                     if arr.grad:
                         grad = self.gradients[arr] = arr.grad
@@ -145,15 +190,21 @@ class Tape:
         return None
 
     def reset(self):
+        """
+        Clear all operations recorded on the tape and zero out all gradients.
+        """
         self.launches = []
         self.zero()
 
     def zero(self):
+        """
+        Zero out all gradients recorded on the tape.
+        """
         for a, g in self.gradients.items():
             if a not in self.const_gradients:
                 if isinstance(a, wp.codegen.StructInstance):
-                    for name in g._struct_.vars:
-                        if isinstance(g._struct_.vars[name].type, wp.array) and g._struct_.vars[name].requires_grad:
+                    for name in g._cls.vars:
+                        if isinstance(g._cls.vars[name].type, wp.array) and g._cls.vars[name].requires_grad:
                             getattr(g, name).zero_()
                 else:
                     g.zero_()

warp/tests/__main__.py

@@ -1,7 +1,4 @@
-from .test_all import run
+from warp.thirdparty.unittest_parallel import main
 
-ret = run()
-
-import sys
-
-sys.exit(ret)
+if __name__ == "__main__":
+    main()