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

warp/builtins.py

@@ -1497,7 +1497,8 @@ add_builtin(
     doc="""Apply the transform to a point ``point`` treating the homogeneous coordinate as w=1.
 
     The transformation is applied treating ``point`` as a column vector, e.g.: ``y = mat*point``.
-    Note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = point^T*mat^T``.
+
+    This is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = point^T*mat^T``.
     If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
     matrix before calling this method.""",
 )
@@ -1515,8 +1516,9 @@ add_builtin(
     group="Vector Math",
     doc="""Apply the transform to a vector ``vec`` treating the homogeneous coordinate as w=0.
 
-    The transformation is applied treating ``vec`` as a column vector, e.g.: ``y = mat*vec``
-    note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = vec^T*mat^T``.
+    The transformation is applied treating ``vec`` as a column vector, e.g.: ``y = mat*vec``.
+
+    This is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = vec^T*mat^T``.
     If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
     matrix before calling this method.""",
 )
@@ -3167,7 +3169,8 @@ def address_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, A
 for array_type in array_types:
     add_builtin(
         "address",
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int},
+        constraint=sametypes,
         defaults={"j": None, "k": None, "l": None},
         hidden=True,
         value_func=address_value_func,
@@ -3211,8 +3214,9 @@ def view_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]
 for array_type in array_types:
     add_builtin(
         "view",
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int},
         defaults={"j": None, "k": None},
+        constraint=sametypes,
         hidden=True,
         value_func=view_value_func,
         group="Utility",
@@ -3254,7 +3258,8 @@ def array_store_value_func(arg_types: Mapping[str, type], arg_values: Mapping[st
 for array_type in array_types:
     add_builtin(
         "array_store",
-        input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=sametypes,
         hidden=True,
         value_func=array_store_value_func,
         skip_replay=True,
@@ -3262,7 +3267,8 @@ for array_type in array_types:
     )
     add_builtin(
         "array_store",
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=sametypes,
         hidden=True,
         value_func=array_store_value_func,
         skip_replay=True,
@@ -3270,7 +3276,8 @@ for array_type in array_types:
     )
     add_builtin(
         "array_store",
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=sametypes,
         hidden=True,
         value_func=array_store_value_func,
         skip_replay=True,
@@ -3278,7 +3285,8 @@ for array_type in array_types:
     )
     add_builtin(
         "array_store",
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=sametypes,
         hidden=True,
         value_func=array_store_value_func,
         skip_replay=True,
@@ -3330,6 +3338,11 @@ add_builtin(
 )
 
 
+def atomic_op_constraint(arg_types: Mapping[str, Any]):
+    idx_types = tuple(arg_types[x] for x in "ijkl" if arg_types.get(x, None) is not None)
+    return all(types_equal(idx_types[0], t) for t in idx_types[1:]) and arg_types["arr"].ndim == len(idx_types)
+
+
 def atomic_op_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     if arg_types is None:
         return Any
@@ -3374,7 +3387,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_add",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically add ``value`` onto ``arr[i]`` and return the old value.",
         group="Utility",
@@ -3383,7 +3397,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_add",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically add ``value`` onto ``arr[i,j]`` and return the old value.",
         group="Utility",
@@ -3392,7 +3407,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_add",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value.",
         group="Utility",
@@ -3401,7 +3417,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_add",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value.",
         group="Utility",
@@ -3411,7 +3428,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_sub",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically subtract ``value`` onto ``arr[i]`` and return the old value.",
         group="Utility",
@@ -3420,7 +3438,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_sub",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value.",
         group="Utility",
@@ -3429,7 +3448,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_sub",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value.",
         group="Utility",
@@ -3438,7 +3458,8 @@ for array_type in array_types:
     add_builtin(
         "atomic_sub",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value.",
         group="Utility",
@@ -3448,44 +3469,48 @@ for array_type in array_types:
     add_builtin(
         "atomic_min",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_min",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_min",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_min",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
@@ -3493,44 +3518,48 @@ for array_type in array_types:
     add_builtin(
         "atomic_max",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_max",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_max",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
     add_builtin(
         "atomic_max",
         hidden=hidden,
-        input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
         value_func=atomic_op_value_func,
         doc="""Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
 
-    .. note:: The operation is only atomic on a per-component basis for vectors and matrices.""",
+    The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
     )
@@ -4295,13 +4324,12 @@ add_builtin(
     value_type=Any,
     doc="""Evaluates a static Python expression and replaces it with its result.
 
-    See the `codegen.html#static-expressions <section on code generation>`_ for more details.
+    See the :ref:`code generation guide <static_expressions>` for more details.
 
-    Note:
-        The inner expression must only reference variables that are available from the current scope where the Warp kernel or function containing the expression is defined,
-        which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
-        The return type of the expression must be either a Warp function, a string, or a type that is supported inside Warp kernels and functions
-        (excluding Warp arrays since they cannot be created in a Warp kernel at the moment).""",
+    The inner expression must only reference variables that are available from the current scope where the Warp kernel or function containing the expression is defined,
+    which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
+    The return type of the expression must be either a Warp function, a string, or a type that is supported inside Warp kernels and functions
+    (excluding Warp arrays since they cannot be created in a Warp kernel at the moment).""",
     group="Code Generation",
 )
 

warp/codegen.py

@@ -939,7 +939,9 @@ class Adjoint:
 
         adj.return_var = None  # return type for function or kernel
         adj.loop_symbols = []  # symbols at the start of each loop
-        adj.loop_const_iter_symbols = []  # iteration variables (constant) for static loops
+        adj.loop_const_iter_symbols = (
+            set()
+        )  # constant iteration variables for static loops (mutating them does not raise an error)
 
         # blocks
         adj.blocks = [Block()]
@@ -2000,22 +2002,11 @@ class Adjoint:
         )
         return range_call
 
-    def begin_record_constant_iter_symbols(adj):
-        if len(adj.loop_const_iter_symbols) > 0:
-            adj.loop_const_iter_symbols.append(adj.loop_const_iter_symbols[-1])
-        else:
-            adj.loop_const_iter_symbols.append(set())
-
-    def end_record_constant_iter_symbols(adj):
-        if len(adj.loop_const_iter_symbols) > 0:
-            adj.loop_const_iter_symbols.pop()
-
     def record_constant_iter_symbol(adj, sym):
-        if len(adj.loop_const_iter_symbols) > 0:
-            adj.loop_const_iter_symbols[-1].add(sym)
+        adj.loop_const_iter_symbols.add(sym)
 
     def is_constant_iter_symbol(adj, sym):
-        return len(adj.loop_const_iter_symbols) > 0 and sym in adj.loop_const_iter_symbols[-1]
+        return sym in adj.loop_const_iter_symbols
 
     def emit_For(adj, node):
         # try and unroll simple range() statements that use constant args
@@ -2045,7 +2036,6 @@ class Adjoint:
                 iter = adj.eval(node.iter)
 
             adj.symbols[node.target.id] = adj.begin_for(iter)
-            adj.begin_record_constant_iter_symbols()
 
             # for loops should be side-effect free, here we store a copy
             adj.loop_symbols.append(adj.symbols.copy())
@@ -2056,7 +2046,6 @@ class Adjoint:
 
             adj.materialize_redefinitions(adj.loop_symbols[-1])
             adj.loop_symbols.pop()
-            adj.end_record_constant_iter_symbols()
 
             adj.end_for(iter)
 
@@ -2559,7 +2548,10 @@ class Adjoint:
                     if warp.config.verify_autograd_array_access:
                         target.mark_write(kernel_name=kernel_name, filename=filename, lineno=lineno)
                 else:
-                    print(f"Warning: in-place op {node.op} is not differentiable")
+                    if warp.config.verbose:
+                        print(f"Warning: in-place op {node.op} is not differentiable")
+                    make_new_assign_statement()
+                    return
 
             # TODO
             elif type_is_vector(target_type) or type_is_quaternion(target_type) or type_is_matrix(target_type):

warp/config.py

@@ -7,7 +7,7 @@
 
 from typing import Optional
 
-version: str = "1.4.1"
+version: str = "1.4.2"
 """Warp version string"""
 
 verify_fp: bool = False

warp/context.py

@@ -1249,6 +1249,7 @@ def add_builtin(
                     key,
                     input_types=arg_types,
                     value_type=return_type,
+                    value_func=value_func if return_type is Any else None,
                     export_func=export_func,
                     dispatch_func=dispatch_func,
                     doc=doc,

warp/examples/core/example_dem.py

@@ -199,9 +199,10 @@ class Example:
 
     # creates a grid of particles
     def particle_grid(self, dim_x, dim_y, dim_z, lower, radius, jitter):
+        rng = np.random.default_rng(42)
         points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
         points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
-        points_t = points_t + np.random.rand(*points_t.shape) * radius * jitter
+        points_t = points_t + rng.random(size=points_t.shape) * radius * jitter
 
         return points_t.reshape((-1, 3))
 

warp/examples/core/example_mesh_intersect.py

@@ -98,11 +98,11 @@ class Example:
 
         for _ in range(self.query_count):
             # random offset
-            p = wp.vec3(rng.random(3) * 0.5 - 0.5) * 5.0
+            p = wp.vec3(rng.random(size=3) * 0.5 - 0.5) * 5.0
 
             # random orientation
-            axis = wp.normalize(wp.vec3(rng.random(3) * 0.5 - 0.5))
-            angle = float(np.random.rand(1)[0])
+            axis = wp.normalize(wp.vec3(rng.random(size=3) * 0.5 - 0.5))
+            angle = rng.random()
 
             q = wp.quat_from_axis_angle(wp.normalize(axis), angle)
 

warp/examples/optim/example_walker.py

@@ -21,7 +21,7 @@ import math
 import os
 
 import numpy as np
-from pxr import Usd, UsdGeom
+from pxr import Gf, Usd, UsdGeom
 
 import warp as wp
 import warp.examples
@@ -93,7 +93,7 @@ class Example:
         geom = UsdGeom.Mesh(asset_stage.GetPrimAtPath("/root/bear"))
         points = geom.GetPointsAttr().Get()
 
-        xform = geom.ComputeLocalToWorldTransform(0.0)
+        xform = Gf.Matrix4f(geom.ComputeLocalToWorldTransform(0.0))
         for i in range(len(points)):
             points[i] = xform.Transform(points[i])
 

warp/examples/sim/example_jacobian_ik.py

@@ -41,6 +41,8 @@ def compute_endeffector_position(
 
 class Example:
     def __init__(self, stage_path="example_jacobian_ik.usd", num_envs=10):
+        rng = np.random.default_rng(42)
+
         builder = wp.sim.ModelBuilder()
 
         self.num_envs = num_envs
@@ -81,7 +83,7 @@ class Example:
             )
             self.target_origin.append((i * 2.0, 4.0, 0.0))
             # joint initial positions
-            builder.joint_q[-3:] = np.random.uniform(-0.5, 0.5, size=3)
+            builder.joint_q[-3:] = rng.uniform(-0.5, 0.5, size=3)
         self.target_origin = np.array(self.target_origin)
 
         # finalize model
@@ -207,6 +209,8 @@ if __name__ == "__main__":
 
     args = parser.parse_known_args()[0]
 
+    rng = np.random.default_rng(42)
+
     with wp.ScopedDevice(args.device):
         example = Example(stage_path=args.stage_path, num_envs=args.num_envs)
 
@@ -218,7 +222,7 @@ if __name__ == "__main__":
         for _ in range(args.num_rollouts):
             # select new random target points for all envs
             example.targets = example.target_origin.copy()
-            example.targets[:, 1:] += np.random.uniform(-0.5, 0.5, size=(example.num_envs, 2))
+            example.targets[:, 1:] += rng.uniform(-0.5, 0.5, size=(example.num_envs, 2))
 
             for iter in range(args.train_iters):
                 example.step()