warp-lang 1.4.0__py3-none-win_amd64.whl → 1.4.2__py3-none-win_amd64.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.""",
 )
@@ -2758,7 +2760,6 @@ add_builtin(
     "rand_init",
     input_types={"seed": int},
     value_type=uint32,
-    export=False,
     group="Random",
     doc="Initialize a new random number generator given a user-defined seed. Returns a 32-bit integer representing the RNG state.",
 )
@@ -2767,7 +2768,6 @@ add_builtin(
     "rand_init",
     input_types={"seed": int, "offset": int},
     value_type=uint32,
-    export=False,
     group="Random",
     doc="""Initialize a new random number generator given a user-defined seed and an offset.
 
@@ -2779,7 +2779,6 @@ add_builtin(
     "randi",
     input_types={"state": uint32},
     value_type=int,
-    export=False,
     group="Random",
     doc="Return a random integer in the range [0, 2^32).",
 )
@@ -2787,7 +2786,6 @@ add_builtin(
     "randi",
     input_types={"state": uint32, "low": int, "high": int},
     value_type=int,
-    export=False,
     group="Random",
     doc="Return a random integer between [low, high).",
 )
@@ -2795,7 +2793,6 @@ add_builtin(
     "randf",
     input_types={"state": uint32},
     value_type=float,
-    export=False,
     group="Random",
     doc="Return a random float between [0.0, 1.0).",
 )
@@ -2803,24 +2800,17 @@ add_builtin(
     "randf",
     input_types={"state": uint32, "low": float, "high": float},
     value_type=float,
-    export=False,
     group="Random",
     doc="Return a random float between [low, high).",
 )
 add_builtin(
-    "randn",
-    input_types={"state": uint32},
-    value_type=float,
-    export=False,
-    group="Random",
-    doc="Sample a normal distribution.",
+    "randn", input_types={"state": uint32}, value_type=float, group="Random", doc="Sample a normal distribution."
 )
 
 add_builtin(
     "sample_cdf",
     input_types={"state": uint32, "cdf": array(dtype=float)},
     value_type=int,
-    export=False,
     group="Random",
     doc="Inverse-transform sample a cumulative distribution function.",
 )
@@ -2828,7 +2818,6 @@ add_builtin(
     "sample_triangle",
     input_types={"state": uint32},
     value_type=vec2,
-    export=False,
     group="Random",
     doc="Uniformly sample a triangle. Returns sample barycentric coordinates.",
 )
@@ -2836,7 +2825,6 @@ add_builtin(
     "sample_unit_ring",
     input_types={"state": uint32},
     value_type=vec2,
-    export=False,
     group="Random",
     doc="Uniformly sample a ring in the xy plane.",
 )
@@ -2844,7 +2832,6 @@ add_builtin(
     "sample_unit_disk",
     input_types={"state": uint32},
     value_type=vec2,
-    export=False,
     group="Random",
     doc="Uniformly sample a disk in the xy plane.",
 )
@@ -2852,7 +2839,6 @@ add_builtin(
     "sample_unit_sphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit sphere surface.",
 )
@@ -2860,7 +2846,6 @@ add_builtin(
     "sample_unit_sphere",
     input_types={"state": uint32},
     value_type=vec3,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit sphere.",
 )
@@ -2868,7 +2853,6 @@ add_builtin(
     "sample_unit_hemisphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit hemisphere surface.",
 )
@@ -2876,7 +2860,6 @@ add_builtin(
     "sample_unit_hemisphere",
     input_types={"state": uint32},
     value_type=vec3,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit hemisphere.",
 )
@@ -2884,7 +2867,6 @@ add_builtin(
     "sample_unit_square",
     input_types={"state": uint32},
     value_type=vec2,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit square.",
 )
@@ -2892,7 +2874,6 @@ add_builtin(
     "sample_unit_cube",
     input_types={"state": uint32},
     value_type=vec3,
-    export=False,
     group="Random",
     doc="Uniformly sample a unit cube.",
 )
@@ -2901,7 +2882,6 @@ add_builtin(
     "poisson",
     input_types={"state": uint32, "lam": float},
     value_type=uint32,
-    export=False,
     group="Random",
     doc="""Generate a random sample from a Poisson distribution.
 
@@ -3189,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,
@@ -3233,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",
@@ -3276,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,
@@ -3284,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,
@@ -3292,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,
@@ -3300,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,
@@ -3352,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
@@ -3396,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",
@@ -3405,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",
@@ -3414,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",
@@ -3423,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",
@@ -3433,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",
@@ -3442,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",
@@ -3451,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",
@@ -3460,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",
@@ -3470,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,
     )
@@ -3515,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,
     )
@@ -4317,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

@@ -777,6 +777,9 @@ def func_match_args(func, arg_types, kwarg_types):
 
 
 def get_arg_type(arg: Union[Var, Any]):
+    if isinstance(arg, str):
+        return str
+
     if isinstance(arg, Sequence):
         return tuple(get_arg_type(x) for x in arg)
 
@@ -936,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 = set()  # 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()]
@@ -1846,7 +1851,7 @@ class Adjoint:
         # detect symbols with conflicting definitions (assigned inside the for loop)
         for items in symbols.items():
             sym = items[0]
-            if adj.loop_const_iter_symbols is not None and sym in adj.loop_const_iter_symbols:
+            if adj.is_constant_iter_symbol(sym):
                 # ignore constant overwriting in for-loops if it is a loop iterator
                 # (it is no problem to unroll static loops multiple times in sequence)
                 continue
@@ -1997,12 +2002,11 @@ class Adjoint:
         )
         return range_call
 
-    def begin_record_constant_iter_symbols(adj):
-        if adj.loop_const_iter_symbols is None:
-            adj.loop_const_iter_symbols = set()
+    def record_constant_iter_symbol(adj, sym):
+        adj.loop_const_iter_symbols.add(sym)
 
-    def end_record_constant_iter_symbols(adj):
-        adj.loop_const_iter_symbols = None
+    def is_constant_iter_symbol(adj, sym):
+        return sym in adj.loop_const_iter_symbols
 
     def emit_For(adj, node):
         # try and unroll simple range() statements that use constant args
@@ -2010,9 +2014,8 @@ class Adjoint:
 
         if isinstance(unroll_range, range):
             const_iter_sym = node.target.id
-            if adj.loop_const_iter_symbols is not None:
-                # prevent constant conflicts in `materialize_redefinitions()`
-                adj.loop_const_iter_symbols.add(const_iter_sym)
+            # prevent constant conflicts in `materialize_redefinitions()`
+            adj.record_constant_iter_symbol(const_iter_sym)
 
             # unroll static for-loop
             for i in unroll_range:
@@ -2033,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())
@@ -2044,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)
 
@@ -2547,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.0"
+version: str = "1.4.2"
 """Warp version string"""
 
 verify_fp: bool = False

warp/context.py

@@ -6,7 +6,6 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import ast
-import builtins
 import ctypes
 import functools
 import hashlib
@@ -22,7 +21,6 @@ import typing
 import weakref
 from copy import copy as shallowcopy
 from pathlib import Path
-from struct import pack as struct_pack
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -1251,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,
@@ -1487,30 +1486,16 @@ class ModuleHasher:
         # hash referenced constants
         for name, value in constants.items():
             ch.update(bytes(name, "utf-8"))
-            # hash the referenced object
-            if isinstance(value, builtins.bool):
-                # This needs to come before the check for `int` since all boolean
-                # values are also instances of `int`.
-                ch.update(struct_pack("?", value))
-            elif isinstance(value, int):
-                ch.update(struct_pack("<q", value))
-            elif isinstance(value, float):
-                ch.update(struct_pack("<d", value))
-            elif isinstance(value, warp.types.float16):
-                # float16 is a special case
-                p = ctypes.pointer(ctypes.c_float(value.value))
-                ch.update(p.contents)
-            elif isinstance(value, tuple(warp.types.scalar_types)):
-                p = ctypes.pointer(value._type_(value.value))
-                ch.update(p.contents)
-            elif isinstance(value, ctypes.Array):
-                ch.update(bytes(value))
-            else:
-                raise RuntimeError(f"Invalid constant type: {type(value)}")
+            ch.update(self.get_constant_bytes(value))
 
         # hash wp.static() expressions that were evaluated at declaration time
         for k, v in adj.static_expressions.items():
-            ch.update(bytes(f"{k} = {v}", "utf-8"))
+            ch.update(bytes(k, "utf-8"))
+            if isinstance(v, Function):
+                if v not in self.functions_in_progress:
+                    ch.update(self.hash_function(v))
+            else:
+                ch.update(self.get_constant_bytes(v))
 
         # hash referenced types
         for t in types.keys():
@@ -1523,6 +1508,24 @@ class ModuleHasher:
 
         return ch.digest()
 
+    def get_constant_bytes(self, value):
+        if isinstance(value, int):
+            # this also handles builtins.bool
+            return bytes(ctypes.c_int(value))
+        elif isinstance(value, float):
+            return bytes(ctypes.c_float(value))
+        elif isinstance(value, warp.types.float16):
+            # float16 is a special case
+            return bytes(ctypes.c_float(value.value))
+        elif isinstance(value, tuple(warp.types.scalar_and_bool_types)):
+            return bytes(value._type_(value.value))
+        elif hasattr(value, "_wp_scalar_type_"):
+            return bytes(value)
+        elif isinstance(value, warp.codegen.StructInstance):
+            return bytes(value._ctype)
+        else:
+            raise TypeError(f"Invalid constant type: {type(value)}")
+
     def get_module_hash(self):
         return self.module_hash
 

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)