warp-lang 1.9.0__py3-none-macosx_10_13_universal2.whl → 1.9.1__py3-none-macosx_10_13_universal2.whl

warp/builtins.py

@@ -126,6 +126,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Scalar),
     doc="Return -1 if ``x`` < 0, return 1 otherwise.",
     group="Scalar Math",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -134,6 +135,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Scalar),
     doc="Return 1.0 if ``x`` < 0.0, return 0.0 otherwise.",
     group="Scalar Math",
+    missing_grad=True,
 )
 add_builtin(
     "nonzero",
@@ -141,6 +143,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Scalar),
     doc="Return 1.0 if ``x`` is not equal to zero, return 0.0 otherwise.",
     group="Scalar Math",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -292,6 +295,7 @@ add_builtin(
 
     This is the most intuitive form of rounding in the colloquial sense, but can be slower than other options like :func:`warp.rint()`.
     Differs from :func:`numpy.round()`, which behaves the same way as :func:`numpy.rint()`.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -302,6 +306,7 @@ add_builtin(
     doc="""Return the nearest integer value to ``x``, rounding halfway cases to nearest even integer.
 
     It is generally faster than :func:`warp.round()`. Equivalent to :func:`numpy.rint()`.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -314,6 +319,7 @@ add_builtin(
     In other words, it discards the fractional part of ``x``.
     It is similar to casting ``float(int(a))``, but preserves the negative sign when ``x`` is in the range [-0.0, -1.0).
     Equivalent to :func:`numpy.trunc()` and :func:`numpy.fix()`.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -322,6 +328,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Float),
     group="Scalar Math",
     doc="""Return the largest integer that is less than or equal to ``x``.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -330,6 +337,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Float),
     group="Scalar Math",
     doc="""Return the smallest integer that is greater than or equal to ``x``.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -340,6 +348,7 @@ add_builtin(
     doc="""Retrieve the fractional part of ``x``.
 
     In other words, it discards the integer part of ``x`` and is equivalent to ``x - trunc(x)``.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -348,6 +357,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Scalar Math",
     doc="""Return ``True`` if ``a`` is a finite number, otherwise return ``False``.""",
+    missing_grad=True,
 )
 add_builtin(
     "isfinite",
@@ -355,6 +365,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if all elements of the vector ``a`` are finite, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isfinite",
@@ -362,6 +373,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if all elements of the quaternion ``a`` are finite, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isfinite",
@@ -369,6 +381,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if all elements of the matrix ``a`` are finite, otherwise return ``False``.",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -377,6 +390,7 @@ add_builtin(
     value_type=builtins.bool,
     doc="Return ``True`` if ``a`` is NaN, otherwise return ``False``.",
     group="Scalar Math",
+    missing_grad=True,
 )
 add_builtin(
     "isnan",
@@ -384,6 +398,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the vector ``a`` is NaN, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isnan",
@@ -391,6 +406,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the quaternion ``a`` is NaN, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isnan",
@@ -398,6 +414,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the matrix ``a`` is NaN, otherwise return ``False``.",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -406,6 +423,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Scalar Math",
     doc="""Return ``True`` if ``a`` is positive or negative infinity, otherwise return ``False``.""",
+    missing_grad=True,
 )
 add_builtin(
     "isinf",
@@ -413,6 +431,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the vector ``a`` is positive or negative infinity, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isinf",
@@ -420,6 +439,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the quaternion ``a`` is positive or negative infinity, otherwise return ``False``.",
+    missing_grad=True,
 )
 add_builtin(
     "isinf",
@@ -427,6 +447,7 @@ add_builtin(
     value_type=builtins.bool,
     group="Vector Math",
     doc="Return ``True`` if any element of the matrix ``a`` is positive or negative infinity, otherwise return ``False``.",
+    missing_grad=True,
 )
 
 
@@ -1182,6 +1203,7 @@ add_builtin(
     doc="Create an identity matrix with shape=(n,n) with the type given by ``dtype``.",
     group="Vector Math",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -1546,6 +1568,7 @@ add_builtin(
     group="Quaternion Math",
     doc="Construct an identity quaternion with zero imaginary part and real part of 1.0",
     export=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -1761,6 +1784,7 @@ add_builtin(
     doc="Construct a spatial transform vector of given dtype.",
     group="Spatial Math",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -1795,6 +1819,7 @@ add_builtin(
     group="Transformations",
     doc="Construct an identity transform with zero translation and identity rotation.",
     export=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -3986,6 +4011,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4039,6 +4065,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4092,6 +4119,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4144,6 +4172,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4196,6 +4225,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4253,6 +4283,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4316,6 +4347,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4379,6 +4411,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -4632,6 +4665,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -4647,6 +4681,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -4656,6 +4691,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -4707,6 +4743,7 @@ add_builtin(
     :param low: The lower bound of the bounding box in BVH space
     :param high: The upper bound of the bounding box in BVH space""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -4722,6 +4759,7 @@ add_builtin(
     :param start: The start of the ray in BVH space
     :param dir: The direction of the ray in BVH space""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -4732,6 +4770,7 @@ add_builtin(
     doc="""Move to the next bound returned by the query.
     The index of the current bound is stored in ``index``, returns ``False`` if there are no more overlapping bound.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5066,12 +5105,13 @@ add_builtin(
     group="Geometry",
     doc="""Construct an axis-aligned bounding box query against a :class:`Mesh`.
 
-    This query can be used to iterate over all triangles inside a volume.
+    This query can be used to iterate over all bounding boxes of the triangles inside a volume.
 
     :param id: The mesh identifier
     :param low: The lower bound of the bounding box in mesh space
     :param high: The upper bound of the bounding box in mesh space""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5079,10 +5119,11 @@ add_builtin(
     input_types={"query": MeshQueryAABB, "index": int},
     value_type=builtins.bool,
     group="Geometry",
-    doc="""Move to the next triangle overlapping the query bounding box.
+    doc="""Move to the next triangle whose bounding box overlaps the query bounding box.
 
     The index of the current face is stored in ``index``, returns ``False`` if there are no more overlapping triangles.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5112,6 +5153,7 @@ add_builtin(
 
     This query can be used to iterate over all neighboring point within a fixed radius from the query point.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5123,6 +5165,7 @@ add_builtin(
 
     The index of the current neighbor is stored in ``index``, returns ``False`` if there are no more neighbors.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5136,6 +5179,7 @@ add_builtin(
 
     Returns -1 if the :class:`HashGrid` has not been reserved.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5147,6 +5191,7 @@ add_builtin(
 
     Returns > 0 if triangles intersect.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5166,6 +5211,7 @@ add_builtin(
     group="Geometry",
     doc="""Evaluates the face normal the mesh given a face index.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5175,6 +5221,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the point of the mesh given a index.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5184,6 +5231,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the velocity of the mesh given a index.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5193,6 +5241,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the point-index of the mesh given a face-vertex index.""",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -5233,12 +5282,32 @@ add_builtin(
 # ---------------------------------
 # Iterators
 
-add_builtin("iter_next", input_types={"range": range_t}, value_type=int, group="Utility", export=False, hidden=True)
 add_builtin(
-    "iter_next", input_types={"query": HashGridQuery}, value_type=int, group="Utility", export=False, hidden=True
+    "iter_next",
+    input_types={"range": range_t},
+    value_type=int,
+    group="Utility",
+    export=False,
+    hidden=True,
+    missing_grad=True,
+)
+add_builtin(
+    "iter_next",
+    input_types={"query": HashGridQuery},
+    value_type=int,
+    group="Utility",
+    export=False,
+    hidden=True,
+    missing_grad=True,
 )
 add_builtin(
-    "iter_next", input_types={"query": MeshQueryAABB}, value_type=int, group="Utility", export=False, hidden=True
+    "iter_next",
+    input_types={"query": MeshQueryAABB},
+    value_type=int,
+    group="Utility",
+    export=False,
+    hidden=True,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5249,6 +5318,7 @@ add_builtin(
     group="Utility",
     doc="""Returns the range in reversed order.""",
     export=False,
+    missing_grad=True,
 )
 
 # ---------------------------------
@@ -5397,6 +5467,7 @@ add_builtin(
     doc="""Returns the value of voxel with coordinates ``i``, ``j``, ``k`` for a volume of type type `dtype`.
 
     If the voxel at this index does not exist, this function returns the background value.""",
+    missing_grad=True,
 )
 
 
@@ -5417,6 +5488,7 @@ add_builtin(
     export=False,
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5447,6 +5519,7 @@ add_builtin(
     doc="""Returns the value of voxel with coordinates ``i``, ``j``, ``k``.
 
     If the voxel at this index does not exist, this function returns the background value""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5455,6 +5528,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5475,6 +5549,7 @@ add_builtin(
     doc="""Returns the vector value of voxel with coordinates ``i``, ``j``, ``k``.
 
     If the voxel at this index does not exist, this function returns the background value.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5483,6 +5558,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5501,6 +5577,7 @@ add_builtin(
     doc="""Returns the :class:`int32` value of voxel with coordinates ``i``, ``j``, ``k``.
 
     If the voxel at this index does not exist, this function returns the background value.""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5509,6 +5586,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -5590,6 +5668,7 @@ add_builtin(
     If the voxel at this index does not exist, this function returns -1.
     This function is available for both index grids and classical volumes.
     """,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5631,6 +5710,7 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="Initialize a new random number generator given a user-defined seed. Returns a 32-bit integer representing the RNG state.",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5642,6 +5722,7 @@ add_builtin(
 
     This alternative constructor can be useful in parallel programs, where a kernel as a whole should share a seed,
     but each thread should generate uncorrelated values. In this case usage should be ``r = rand_init(seed, tid)``""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5650,6 +5731,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Return a random integer in the range [-2^31, 2^31).",
+    missing_grad=True,
 )
 add_builtin(
     "randi",
@@ -5657,6 +5739,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Return a random integer between [low, high).",
+    missing_grad=True,
 )
 add_builtin(
     "randu",
@@ -5664,6 +5747,7 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="Return a random unsigned integer in the range [0, 2^32).",
+    missing_grad=True,
 )
 add_builtin(
     "randu",
@@ -5671,6 +5755,7 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="Return a random unsigned integer between [low, high).",
+    missing_grad=True,
 )
 add_builtin(
     "randf",
@@ -5678,6 +5763,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Return a random float between [0.0, 1.0).",
+    missing_grad=True,
 )
 add_builtin(
     "randf",
@@ -5685,6 +5771,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Return a random float between [low, high).",
+    missing_grad=True,
 )
 add_builtin(
     "randn",
@@ -5692,6 +5779,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Sample a normal (Gaussian) distribution of mean 0 and variance 1. ",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5700,6 +5788,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Inverse-transform sample a cumulative distribution function.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_triangle",
@@ -5707,6 +5796,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a triangle. Returns sample barycentric coordinates.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_ring",
@@ -5714,6 +5804,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a ring in the xy plane.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_disk",
@@ -5721,6 +5812,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a disk in the xy plane.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_sphere_surface",
@@ -5728,6 +5820,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit sphere surface.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_sphere",
@@ -5735,6 +5828,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit sphere.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_hemisphere_surface",
@@ -5742,6 +5836,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit hemisphere surface.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_hemisphere",
@@ -5749,6 +5844,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit hemisphere.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_square",
@@ -5756,6 +5852,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a unit square.",
+    missing_grad=True,
 )
 add_builtin(
     "sample_unit_cube",
@@ -5763,6 +5860,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit cube.",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5774,6 +5872,7 @@ add_builtin(
 
     :param state: RNG state
     :param lam: The expected value of the distribution""",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5891,9 +5990,16 @@ add_builtin(
     dispatch_func=printf_dispatch_func,
     group="Utility",
     doc="Allows printing formatted strings using C-style format specifiers.",
+    missing_grad=True,
 )
 
-add_builtin("print", input_types={"value": Any}, doc="Print variable to stdout", export=False, group="Utility")
+add_builtin(
+    "print",
+    input_types={"value": Any},
+    doc="Print variable to stdout",
+    export=False,
+    group="Utility",
+)
 
 add_builtin(
     "breakpoint",
@@ -5903,6 +6009,7 @@ add_builtin(
     group="Utility",
     namespace="",
     native_func="__debugbreak",
+    missing_grad=True,
 )
 
 # helpers
@@ -5920,6 +6027,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid1d",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5930,6 +6038,7 @@ add_builtin(
     doc="Returns the number of threads in the current block.",
     namespace="",
     native_func="builtin_block_dim",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5944,6 +6053,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid2d",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5958,6 +6068,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid3d",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -5972,17 +6083,37 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid4d",
+    missing_grad=True,
 )
 
 
+def copy_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    a = arg_types["a"]
+
+    # if the input is a shared tile, we force a copy
+    if is_tile(a) and a.storage == "shared":
+        return tile(
+            dtype=a.dtype,
+            shape=a.shape,
+            storage=a.storage,
+            strides=a.strides,
+            layout=a.layout,
+            owner=True,
+        )
+
+    return a
+
+
 add_builtin(
     "copy",
     input_types={"a": Any},
-    value_func=lambda arg_types, arg_values: arg_types["a"],
+    value_func=copy_value_func,
     hidden=True,
     export=False,
     group="Utility",
 )
+
+
 add_builtin(
     "assign",
     input_types={"dest": Any, "src": Any},
@@ -5992,6 +6123,37 @@ add_builtin(
 )
 
 
+def select_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    if arg_types is None:
+        return Any
+
+    v_true = arg_types["value_if_true"]
+    v_false = arg_types["value_if_false"]
+
+    if not types_equal(v_true, v_false):
+        raise RuntimeError(
+            f"select() true value type ({v_true}) must be of the same type as the false type ({v_false})"
+        )
+
+    if is_tile(v_false):
+        if v_true.storage == "register":
+            return v_true
+        if v_false.storage == "register":
+            return v_false
+
+        # both v_true and v_false are shared
+        return tile(
+            dtype=v_true.dtype,
+            shape=v_true.shape,
+            storage=v_true.storage,
+            strides=v_true.strides,
+            layout=v_true.layout,
+            owner=True,
+        )
+
+    return v_true
+
+
 def select_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
     warp.utils.warn(
         "wp.select() is deprecated and will be removed in a future\n"
@@ -6008,7 +6170,7 @@ def select_dispatch_func(input_types: Mapping[str, type], return_type: Any, args
 add_builtin(
     "select",
     input_types={"cond": builtins.bool, "value_if_false": Any, "value_if_true": Any},
-    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
+    value_func=select_value_func,
     dispatch_func=select_dispatch_func,
     doc="""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
 
@@ -6021,7 +6183,7 @@ 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: Any if arg_types is None else arg_types["value_if_false"],
+        value_func=select_value_func,
         dispatch_func=select_dispatch_func,
         doc="""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
 
@@ -6033,7 +6195,7 @@ for t in int_types:
 add_builtin(
     "select",
     input_types={"arr": array(dtype=Any), "value_if_false": Any, "value_if_true": Any},
-    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
+    value_func=select_value_func,
     dispatch_func=select_dispatch_func,
     doc="""Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``.
 
@@ -6043,10 +6205,40 @@ add_builtin(
     group="Utility",
 )
 
+
+def where_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    if arg_types is None:
+        return Any
+
+    v_true = arg_types["value_if_true"]
+    v_false = arg_types["value_if_false"]
+
+    if not types_equal(v_true, v_false):
+        raise RuntimeError(f"where() true value type ({v_true}) must be of the same type as the false type ({v_false})")
+
+    if is_tile(v_false):
+        if v_true.storage == "register":
+            return v_true
+        if v_false.storage == "register":
+            return v_false
+
+        # both v_true and v_false are shared
+        return tile(
+            dtype=v_true.dtype,
+            shape=v_true.shape,
+            storage=v_true.storage,
+            strides=v_true.strides,
+            layout=v_true.layout,
+            owner=True,
+        )
+
+    return v_true
+
+
 add_builtin(
     "where",
     input_types={"cond": builtins.bool, "value_if_true": Any, "value_if_false": Any},
-    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
+    value_func=where_value_func,
     doc="Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``.",
     group="Utility",
 )
@@ -6054,14 +6246,14 @@ for t in int_types:
     add_builtin(
         "where",
         input_types={"cond": t, "value_if_true": Any, "value_if_false": Any},
-        value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
+        value_func=where_value_func,
         doc="Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``.",
         group="Utility",
     )
 add_builtin(
     "where",
     input_types={"arr": array(dtype=Any), "value_if_true": Any, "value_if_false": Any},
-    value_func=lambda arg_types, arg_values: Any if arg_types is None else arg_types["value_if_false"],
+    value_func=where_value_func,
     doc="Select between two arguments, if ``arr`` is not null then return ``value_if_true``, otherwise return ``value_if_false``.",
     group="Utility",
 )
@@ -6321,6 +6513,7 @@ add_builtin(
     hidden=True,
     skip_replay=True,
     group="Utility",
+    missing_grad=True,
 )
 
 
@@ -6337,6 +6530,7 @@ add_builtin(
     dispatch_func=load_dispatch_func,
     hidden=True,
     group="Utility",
+    missing_grad=True,
 )
 
 
@@ -6653,6 +6847,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_cas",
@@ -6666,6 +6861,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_cas",
@@ -6679,6 +6875,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_cas",
@@ -6700,6 +6897,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
 
     add_builtin(
@@ -6714,6 +6912,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_exch",
@@ -6727,6 +6926,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_exch",
@@ -6740,6 +6940,7 @@ for array_type in array_types:
     The operation is only atomic on a per-component basis for vectors and matrices.""",
         group="Utility",
         skip_replay=True,
+        missing_grad=True,
     )
     add_builtin(
         "atomic_exch",
@@ -6903,6 +7104,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &quaternion[index]
 add_builtin(
@@ -6913,6 +7115,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &transformation[index]
 add_builtin(
@@ -6923,6 +7126,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &(*vector)[index]
 add_builtin(
@@ -6933,6 +7137,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &(*matrix)[i, j]
 add_builtin(
@@ -6943,6 +7148,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &(*quaternion)[index]
 add_builtin(
@@ -6953,6 +7159,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 # implements &(*transformation)[index]
 add_builtin(
@@ -6963,6 +7170,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 
 
@@ -7173,6 +7381,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 
 
@@ -7191,6 +7400,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
+    missing_grad=True,
 )
 
 
@@ -7401,6 +7611,7 @@ for t in scalar_types + vector_types + (bool,):
         doc="Prints an error to stdout if ``a`` and ``b`` are not equal",
         group="Utility",
         hidden=True,
+        missing_grad=True,
     )
 
     add_builtin(
@@ -7411,6 +7622,7 @@ for t in scalar_types + vector_types + (bool,):
         group="Utility",
         hidden=True,
         export=False,
+        missing_grad=True,
     )
 
 
@@ -7429,6 +7641,7 @@ add_builtin(
     doc="Prints an error to stdout if ``a`` and ``b`` are not equal",
     group="Utility",
     hidden=True,
+    missing_grad=True,
 )
 add_builtin(
     "expect_neq",
@@ -7439,6 +7652,7 @@ add_builtin(
     group="Utility",
     hidden=True,
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -7449,6 +7663,7 @@ add_builtin(
     doc="Prints an error to stdout if ``a`` and ``b`` are not equal",
     group="Utility",
     hidden=True,
+    missing_grad=True,
 )
 add_builtin(
     "expect_neq",
@@ -7459,6 +7674,7 @@ add_builtin(
     group="Utility",
     hidden=True,
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -7549,6 +7765,7 @@ add_builtin(
     value_type=None,
     doc="Prints an error to stdout if ``a`` and ``b`` are not closer than tolerance in magnitude",
     group="Utility",
+    missing_grad=True,
 )
 add_builtin(
     "expect_near",
@@ -7558,6 +7775,7 @@ add_builtin(
     value_type=None,
     doc="Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude",
     group="Utility",
+    missing_grad=True,
 )
 add_builtin(
     "expect_near",
@@ -7567,6 +7785,7 @@ add_builtin(
     value_type=None,
     doc="Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude",
     group="Utility",
+    missing_grad=True,
 )
 add_builtin(
     "expect_near",
@@ -7580,6 +7799,7 @@ add_builtin(
     value_type=None,
     doc="Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude",
     group="Utility",
+    missing_grad=True,
 )
 
 # ---------------------------------
@@ -7590,6 +7810,7 @@ add_builtin(
     input_types={"arr": array(dtype=Scalar), "value": Scalar},
     value_type=int,
     doc="Search a sorted array ``arr`` for the closest element greater than or equal to ``value``.",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -7597,6 +7818,7 @@ add_builtin(
     input_types={"arr": array(dtype=Scalar), "arr_begin": int, "arr_end": int, "value": Scalar},
     value_type=int,
     doc="Search a sorted array ``arr`` in the range [arr_begin, arr_end) for the closest element greater than or equal to ``value``.",
+    missing_grad=True,
 )
 
 # ---------------------------------
@@ -7672,13 +7894,36 @@ add_builtin(
 )
 
 # bitwise operators
-add_builtin("bit_and", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int))
-add_builtin("bit_or", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int))
-add_builtin("bit_xor", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int))
-add_builtin("lshift", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int))
-add_builtin("rshift", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int))
-add_builtin("invert", input_types={"a": Int}, value_func=sametypes_create_value_func(Int))
-
+add_builtin(
+    "bit_and",
+    input_types={"a": Int, "b": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    missing_grad=True,
+)
+add_builtin(
+    "bit_or",
+    input_types={"a": Int, "b": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    missing_grad=True,
+)
+add_builtin(
+    "bit_xor",
+    input_types={"a": Int, "b": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    missing_grad=True,
+)
+add_builtin("lshift", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int), group="Operators")
+add_builtin(
+    "rshift",
+    input_types={"a": Int, "b": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    missing_grad=True,
+)
+add_builtin("invert", input_types={"a": Int}, value_func=sametypes_create_value_func(Int), group="Operators")
 
 add_builtin(
     "mul", input_types={"a": Scalar, "b": Scalar}, value_func=sametypes_create_value_func(Scalar), group="Operators"
@@ -7878,6 +8123,7 @@ add_builtin(
     value_func=sametypes_create_value_func(vector(length=Any, dtype=Scalar)),
     doc="Modulo operation using truncated division.",
     group="Operators",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -7937,6 +8183,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Scalar),
     doc="",
     group="Operators",
+    missing_grad=True,
 )
 
 add_builtin("pos", input_types={"x": Scalar}, value_func=sametypes_create_value_func(Scalar), group="Operators")
@@ -7984,12 +8231,16 @@ add_builtin(
     group="Operators",
 )
 
-add_builtin("unot", input_types={"a": builtins.bool}, value_type=builtins.bool, doc="", group="Operators")
+add_builtin(
+    "unot", input_types={"a": builtins.bool}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True
+)
 for t in int_types:
-    add_builtin("unot", input_types={"a": t}, value_type=builtins.bool, doc="", group="Operators")
+    add_builtin("unot", input_types={"a": t}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True)
 
 
-add_builtin("unot", input_types={"a": array(dtype=Any)}, value_type=builtins.bool, doc="", group="Operators")
+add_builtin(
+    "unot", input_types={"a": array(dtype=Any)}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True
+)
 
 
 # Tile operators
@@ -8183,6 +8434,7 @@ add_builtin(
     doc="Add a square matrix and a diagonal matrix 'd' represented as a 1D tile",
     group="Tile Primitives",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -8510,6 +8762,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 add_builtin(
@@ -8531,6 +8784,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 
@@ -8655,6 +8909,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 
@@ -8785,6 +9040,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 
@@ -8917,6 +9173,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 
@@ -9049,6 +9306,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
+    missing_grad=True,
 )
 
 
@@ -9068,6 +9326,7 @@ add_builtin(
     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",
+    missing_grad=True,
 )
 
 
@@ -9092,6 +9351,7 @@ add_builtin(
     doc="Return the number of elements in a vector.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -9101,6 +9361,7 @@ add_builtin(
     doc="Return the number of elements in a quaternion.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -9110,6 +9371,7 @@ add_builtin(
     doc="Return the number of rows in a matrix.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -9119,6 +9381,7 @@ add_builtin(
     doc="Return the number of elements in a transformation.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -9128,6 +9391,7 @@ add_builtin(
     doc="Return the size of the first dimension in an array.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 add_builtin(
@@ -9137,6 +9401,7 @@ add_builtin(
     doc="Return the number of rows in a tile.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 
@@ -9211,6 +9476,7 @@ add_builtin(
     doc="Return the number of elements in a tuple.",
     group="Utility",
     export=False,
+    missing_grad=True,
 )
 
 # ---------------------------------