warp-lang 1.3.1__py3-none-macosx_10_13_universal2.whl → 1.3.2__py3-none-macosx_10_13_universal2.whl

warp/autograd.py

@@ -593,13 +593,13 @@ def jacobian(
         input_output_mask = []
     arg_names = [arg.label for arg in kernel.adj.args]
 
-    def resolve_arg(name):
+    def resolve_arg(name, offset: int = 0):
         if isinstance(name, int):
             return name
-        return arg_names.index(name)
+        return arg_names.index(name) + offset
 
     input_output_mask = [
-        (resolve_arg(input_name), resolve_arg(output_name) - len(inputs))
+        (resolve_arg(input_name), resolve_arg(output_name, -len(inputs)))
         for input_name, output_name in input_output_mask
     ]
     input_output_mask = set(input_output_mask)
@@ -694,13 +694,13 @@ def jacobian_fd(
         input_output_mask = []
     arg_names = [arg.label for arg in kernel.adj.args]
 
-    def resolve_arg(name):
+    def resolve_arg(name, offset: int = 0):
         if isinstance(name, int):
             return name
-        return arg_names.index(name)
+        return arg_names.index(name) + offset
 
     input_output_mask = [
-        (resolve_arg(input_name), resolve_arg(output_name) - len(inputs))
+        (resolve_arg(input_name), resolve_arg(output_name, -len(inputs)))
         for input_name, output_name in input_output_mask
     ]
     input_output_mask = set(input_output_mask)

warp/builtins.py

@@ -549,7 +549,9 @@ add_builtin(
 add_builtin(
     "skew",
     input_types={"vec": vector(length=3, dtype=Scalar)},
-    value_func=lambda arg_types, arg_values: matrix(shape=(3, 3), dtype=arg_types["vec"]._wp_scalar_type_),
+    value_func=lambda arg_types, arg_values: matrix(shape=(3, 3), dtype=Scalar)
+    if arg_types is None
+    else matrix(shape=(3, 3), dtype=arg_types["vec"]._wp_scalar_type_),
     group="Vector Math",
     doc="Compute the skew-symmetric 3x3 matrix for a 3D vector ``vec``.",
 )
@@ -603,9 +605,9 @@ add_builtin(
 add_builtin(
     "transpose",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar)},
-    value_func=lambda arg_types, arg_values: matrix(
-        shape=(arg_types["a"]._shape_[1], arg_types["a"]._shape_[0]), dtype=arg_types["a"]._wp_scalar_type_
-    ),
+    value_func=lambda arg_types, arg_values: matrix(shape=(Any, Any), dtype=Scalar)
+    if arg_types is None
+    else matrix(shape=(arg_types["a"]._shape_[1], arg_types["a"]._shape_[0]), dtype=arg_types["a"]._wp_scalar_type_),
     group="Vector Math",
     doc="Return the transpose of the matrix ``a``.",
 )
@@ -1652,14 +1654,18 @@ add_builtin(
 add_builtin(
     "spatial_top",
     input_types={"svec": vector(length=6, dtype=Float)},
-    value_func=lambda arg_types, arg_values: vector(length=3, dtype=arg_types["svec"]._wp_scalar_type_),
+    value_func=lambda arg_types, arg_values: vector(length=3, dtype=Float)
+    if arg_types is None
+    else vector(length=3, dtype=arg_types["svec"]._wp_scalar_type_),
     group="Spatial Math",
     doc="Return the top (first) part of a 6D screw vector.",
 )
 add_builtin(
     "spatial_bottom",
     input_types={"svec": vector(length=6, dtype=Float)},
-    value_func=lambda arg_types, arg_values: vector(length=3, dtype=arg_types["svec"]._wp_scalar_type_),
+    value_func=lambda arg_types, arg_values: vector(length=3, dtype=Float)
+    if arg_types is None
+    else vector(length=3, dtype=arg_types["svec"]._wp_scalar_type_),
     group="Spatial Math",
     doc="Return the bottom (second) part of a 6D screw vector.",
 )
@@ -3079,7 +3085,7 @@ add_builtin(
 add_builtin(
     "select",
     input_types={"cond": builtins.bool, "value_if_false": Any, "value_if_true": Any},
-    value_func=lambda arg_types, arg_values: arg_types["value_if_false"],
+    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
     doc="Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``",
     group="Utility",
 )
@@ -3087,14 +3093,14 @@ for t in int_types:
     add_builtin(
         "select",
         input_types={"cond": t, "value_if_false": Any, "value_if_true": Any},
-        value_func=lambda arg_types, arg_values: arg_types["value_if_false"],
+        value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
         doc="Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``",
         group="Utility",
     )
 add_builtin(
     "select",
     input_types={"arr": array(dtype=Any), "value_if_false": Any, "value_if_true": Any},
-    value_func=lambda arg_types, arg_values: arg_types["value_if_false"],
+    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
     doc="Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``",
     group="Utility",
 )
@@ -3145,14 +3151,9 @@ def address_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, A
 
     idx_count = len(idx_types)
 
-    if idx_count < arr_type.ndim:
+    if idx_count != arr_type.ndim:
         raise RuntimeError(
-            "Num indices < num dimensions for array load, this is a codegen error, should have generated a view instead"
-        )
-
-    if idx_count > arr_type.ndim:
-        raise RuntimeError(
-            f"Num indices > num dimensions for array load, received {idx_count}, but array only has {arr_type.ndim}"
+            f"The number of indices provided ({idx_count}) does not match the array dimensions ({arr_type.ndim}) for array load"
         )
 
     # check index types
@@ -3229,14 +3230,9 @@ def array_store_value_func(arg_types: Mapping[str, type], arg_values: Mapping[st
 
     idx_count = len(idx_types)
 
-    if idx_count < arr_type.ndim:
-        raise RuntimeError(
-            "Num indices < num dimensions for array store, this is a codegen error, should have generated a view instead"
-        )
-
-    if idx_count > arr_type.ndim:
+    if idx_count != arr_type.ndim:
         raise RuntimeError(
-            f"Num indices > num dimensions for array store, received {idx_count}, but array only has {arr_type.ndim}"
+            f"The number of indices provided ({idx_count}) does not match the array dimensions ({arr_type.ndim}) for array store"
         )
 
     # check index types
@@ -3335,6 +3331,9 @@ add_builtin(
 
 
 def atomic_op_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    if arg_types is None:
+        return Any
+
     arr_type = arg_types["arr"]
     value_type = arg_types["value"]
     idx_types = tuple(arg_types[x] for x in "ijkl" if arg_types.get(x, None) is not None)
@@ -3377,7 +3376,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto ``arr[i]``.",
+        doc="Atomically add ``value`` onto ``arr[i]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3386,7 +3385,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto ``arr[i,j]``.",
+        doc="Atomically add ``value`` onto ``arr[i,j]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3395,7 +3394,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto ``arr[i,j,k]``.",
+        doc="Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3404,7 +3403,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto ``arr[i,j,k,l]``.",
+        doc="Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3414,7 +3413,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto ``arr[i]``.",
+        doc="Atomically subtract ``value`` onto ``arr[i]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3423,7 +3422,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto ``arr[i,j]``.",
+        doc="Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3432,7 +3431,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto ``arr[i,j,k]``.",
+        doc="Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3441,7 +3440,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto ``arr[i,j,k,l]``.",
+        doc="Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value.",
         group="Utility",
         skip_replay=True,
     )
@@ -3451,7 +3450,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the minimum of ``value`` and ``arr[i]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3462,7 +3461,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the minimum of ``value`` and ``arr[i,j]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3473,7 +3472,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the minimum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3484,7 +3483,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the minimum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3496,7 +3495,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the maximum of ``value`` and ``arr[i]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3507,7 +3506,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the maximum of ``value`` and ``arr[i,j]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3518,7 +3517,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the maximum of ``value`` and ``arr[i,j,k]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -3529,7 +3528,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"arr": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="""Compute the maximum of ``value`` and ``arr[i,j,k,l]`` and atomically update the array.
+        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.""",
         group="Utility",
@@ -4056,7 +4055,7 @@ def matmat_mul_constraint(arg_types: Mapping[str, type]):
 
 def matmat_mul_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     if arg_types is None:
-        return matrix(length=Any, dtype=Scalar)
+        return matrix(shape=(Any, Any), dtype=Scalar)
 
     if arg_types["a"]._wp_scalar_type_ != arg_types["b"]._wp_scalar_type_:
         raise RuntimeError(
@@ -4102,7 +4101,11 @@ add_builtin(
 )
 
 add_builtin(
-    "mod", input_types={"a": Scalar, "b": Scalar}, value_func=sametypes_create_value_func(Scalar), group="Operators"
+    "mod",
+    input_types={"a": Scalar, "b": Scalar},
+    value_func=sametypes_create_value_func(Scalar),
+    doc="Modulo operation using truncated division.",
+    group="Operators",
 )
 
 add_builtin(

warp/codegen.py

@@ -795,7 +795,13 @@ class Adjoint:
         # extract name of source file
         adj.filename = inspect.getsourcefile(func) or "unknown source file"
         # get source file line number where function starts
-        _, adj.fun_lineno = inspect.getsourcelines(func)
+        try:
+            _, adj.fun_lineno = inspect.getsourcelines(func)
+        except OSError as e:
+            raise RuntimeError(
+                "Directly evaluating Warp code defined as a string using `exec()` is not supported, "
+                "please save it on a file and use `importlib` if needed."
+            ) from e
 
         # get function source code
         adj.source = inspect.getsource(func)
@@ -1592,15 +1598,7 @@ class Adjoint:
         if node.id in adj.symbols:
             return adj.symbols[node.id]
 
-        # try and resolve the name using the function's globals context (used to lookup constants + functions)
-        obj = adj.func.__globals__.get(node.id)
-
-        if obj is None:
-            # Lookup constant in captured contents
-            capturedvars = dict(
-                zip(adj.func.__code__.co_freevars, [c.cell_contents for c in (adj.func.__closure__ or [])])
-            )
-            obj = capturedvars.get(str(node.id), None)
+        obj = adj.resolve_external_reference(node.id)
 
         if obj is None:
             raise WarpCodegenKeyError("Referencing undefined symbol: " + str(node.id))
@@ -2299,7 +2297,9 @@ class Adjoint:
                     )
 
             else:
-                raise WarpCodegenError("Can only subscript assign array, vector, quaternion, and matrix types")
+                raise WarpCodegenError(
+                    f"Can only subscript assign array, vector, quaternion, and matrix types, got {target_type}"
+                )
 
         elif isinstance(lhs, ast.Name):
             # symbol name
@@ -2448,24 +2448,11 @@ class Adjoint:
         if path[0] in __builtins__:
             return __builtins__[path[0]]
 
-        # Look up the closure info and append it to adj.func.__globals__
-        # in case you want to define a kernel inside a function and refer
-        # to variables you've declared inside that function:
-        def extract_contents(contents):
-            return contents if isinstance(contents, warp.context.Function) or not callable(contents) else contents
-
-        capturedvars = dict(
-            zip(
-                adj.func.__code__.co_freevars, [extract_contents(c.cell_contents) for c in (adj.func.__closure__ or [])]
-            )
-        )
-        vars_dict = {**adj.func.__globals__, **capturedvars}
-
-        if path[0] in vars_dict:
-            expr = vars_dict[path[0]]
+        # look up in closure/global variables
+        expr = adj.resolve_external_reference(path[0])
 
         # Support Warp types in kernels without the module suffix (e.g. v = vec3(0.0,0.2,0.4)):
-        else:
+        if expr is None:
             expr = getattr(warp, path[0], None)
 
         if expr:
@@ -2526,6 +2513,16 @@ class Adjoint:
 
         return None, path
 
+    def resolve_external_reference(adj, name: str):
+        try:
+            # look up in closure variables
+            idx = adj.func.__code__.co_freevars.index(name)
+            obj = adj.func.__closure__[idx].cell_contents
+        except ValueError:
+            # look up in global variables
+            obj = adj.func.__globals__.get(name)
+        return obj
+
     # annotate generated code with the original source code line
     def set_lineno(adj, lineno):
         if adj.lineno is None or adj.lineno != lineno:
@@ -2551,17 +2548,8 @@ class Adjoint:
 
         for node in ast.walk(adj.tree):
             if isinstance(node, ast.Name) and node.id not in local_variables:
-                # This node could be a reference to a wp.constant defined or imported into the current namespace
-
-                # try and resolve the name using the function's globals context (used to lookup constants + functions)
-                obj = adj.func.__globals__.get(node.id)
-
-                if obj is None:
-                    # Lookup constant in captured contents
-                    capturedvars = dict(
-                        zip(adj.func.__code__.co_freevars, [c.cell_contents for c in (adj.func.__closure__ or [])])
-                    )
-                    obj = capturedvars.get(str(node.id), None)
+                # look up in closure/global variables
+                obj = adj.resolve_external_reference(node.id)
 
                 if warp.types.is_value(obj):
                     constants_dict[node.id] = obj
@@ -2571,6 +2559,7 @@ class Adjoint:
 
                 if warp.types.is_value(obj):
                     constants_dict[".".join(path)] = obj
+
             elif isinstance(node, ast.Assign):
                 # Add the LHS names to the local_variables so we know any subsequent uses are shadowed
                 lhs = node.targets[0]

warp/config.py

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