warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0__py3-none-manylinux_2_34_aarch64.whl

warp/{builtins.py → _src/builtins.py}

@@ -20,14 +20,16 @@ import functools
 import math
 from typing import Any, Callable, Mapping, Sequence
 
-import warp.build
-import warp.context
-import warp.utils
-from warp.codegen import Reference, Var, get_arg_value, strip_reference
-from warp.types import *
+import warp._src.build
+import warp._src.context
+import warp._src.utils
+from warp._src.codegen import Reference, Var, get_arg_value, strip_reference
+from warp._src.types import *
 
 from .context import add_builtin
 
+_wp_module_name_ = "warp.builtins"
+
 
 def seq_check_equal(seq_1, seq_2):
     if not isinstance(seq_1, Sequence) or not isinstance(seq_2, Sequence):
@@ -61,11 +63,11 @@ def sametypes_create_value_func(default: TypeVar):
 
 def extract_tuple(arg, as_constant=False):
     if isinstance(arg, Var):
-        if isinstance(arg.type, warp.types.tuple_t):
+        if isinstance(arg.type, warp._src.types.tuple_t):
             out = arg.type.values
         else:
             out = (arg,)
-    elif isinstance(arg, warp.types.tuple_t):
+    elif isinstance(arg, warp._src.types.tuple_t):
         out = arg.values
     elif not isinstance(arg, Sequence):
         out = (arg,)
@@ -82,7 +84,7 @@ def static_len_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str
     if arg_types is None:
         return int
 
-    length = warp.types.type_length(arg_types["a"])
+    length = warp._src.types.type_length(arg_types["a"])
     return Var(None, type=int, constant=length)
 
 
@@ -126,7 +128,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -135,7 +137,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,
+    is_differentiable=False,
 )
 add_builtin(
     "nonzero",
@@ -143,7 +145,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -285,7 +287,36 @@ add_builtin(
     group="Scalar Math",
     require_original_output_arg=True,
 )
-
+add_builtin(
+    "erf",
+    input_types={"x": Float},
+    value_func=sametypes_create_value_func(Float),
+    doc="Return the error function of ``x``.",
+    group="Scalar Math",
+)
+add_builtin(
+    "erfc",
+    input_types={"x": Float},
+    value_func=sametypes_create_value_func(Float),
+    doc="Return the complementary error function of ``x``.",
+    group="Scalar Math",
+)
+add_builtin(
+    "erfinv",
+    input_types={"x": Float},
+    value_func=sametypes_create_value_func(Float),
+    doc="Return the inverse error function of ``x``.",
+    group="Scalar Math",
+    require_original_output_arg=True,
+)
+add_builtin(
+    "erfcinv",
+    input_types={"x": Float},
+    value_func=sametypes_create_value_func(Float),
+    doc="Return the inverse complementary error function of ``x``.",
+    group="Scalar Math",
+    require_original_output_arg=True,
+)
 add_builtin(
     "round",
     input_types={"x": Float},
@@ -295,7 +326,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -306,7 +337,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -319,7 +350,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -328,7 +359,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -337,7 +368,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -348,7 +379,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -357,7 +388,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isfinite",
@@ -365,7 +396,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isfinite",
@@ -373,7 +404,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isfinite",
@@ -381,7 +412,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -390,7 +421,7 @@ add_builtin(
     value_type=builtins.bool,
     doc="Return ``True`` if ``a`` is NaN, otherwise return ``False``.",
     group="Scalar Math",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "isnan",
@@ -398,7 +429,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isnan",
@@ -406,7 +437,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isnan",
@@ -414,7 +445,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -423,7 +454,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isinf",
@@ -431,7 +462,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isinf",
@@ -439,7 +470,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,
+    is_differentiable=False,
 )
 add_builtin(
     "isinf",
@@ -447,7 +478,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,
+    is_differentiable=False,
 )
 
 
@@ -555,7 +586,7 @@ add_builtin(
     value_func=lambda arg_types, arg_values: warp.uint32,
     doc="Return the index of the minimum element of a vector ``a``.",
     group="Vector Math",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "argmax",
@@ -563,7 +594,7 @@ add_builtin(
     value_func=lambda arg_types, arg_values: warp.uint32,
     doc="Return the index of the maximum element of a vector ``a``.",
     group="Vector Math",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -888,7 +919,7 @@ def vector_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
 
             if dtype is None:
                 dtype = value_type
-            elif not warp.types.scalars_equal(value_type, dtype):
+            elif not warp._src.types.scalars_equal(value_type, dtype):
                 raise RuntimeError(
                     f"the value used to fill this vector is expected to be of the type `{dtype.__name__}`"
                 )
@@ -909,7 +940,7 @@ def vector_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
 
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"all values used to initialize this vector are expected to be of the type `{dtype.__name__}`"
             )
@@ -992,7 +1023,7 @@ def matrix_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
 
             if dtype is None:
                 dtype = value_type
-            elif not warp.types.scalars_equal(value_type, dtype):
+            elif not warp._src.types.scalars_equal(value_type, dtype):
                 raise RuntimeError(
                     f"the value used to fill this matrix is expected to be of the type `{dtype.__name__}`"
                 )
@@ -1002,7 +1033,7 @@ def matrix_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
             raise RuntimeError("the `shape` argument must be specified when initializing a matrix by value")
 
         if all(type_is_vector(x) for x in variadic_arg_types):
-            warp.utils.warn(
+            warp._src.utils.warn(
                 "the built-in `wp.matrix()` won't support taking column vectors as input "
                 "in the future. Use `wp.matrix_from_rows()` or `wp.matrix_from_cols()` instead.",
                 DeprecationWarning,
@@ -1031,7 +1062,7 @@ def matrix_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
 
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"all values used to initialize this matrix are expected to be of the type `{dtype.__name__}`"
             )
@@ -1203,49 +1234,18 @@ 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,
+    is_differentiable=False,
 )
 
 
 def matrix_transform_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
-    warp.utils.warn(
-        "the built-in `wp.matrix()` function to construct a 4x4 matrix from a 3D position, quaternion, "
-        "and 3D scale vector will be deprecated in favor of `wp.transform_compose()`.",
-        DeprecationWarning,
-    )
     if arg_types is None:
         return matrix(shape=(4, 4), dtype=Float)
 
-    dtype = arg_values.get("dtype", None)
-
-    value_arg_types = tuple(v for k, v in arg_types.items() if k != "dtype")
-    try:
-        value_type = scalar_infer_type(value_arg_types)
-    except RuntimeError:
-        raise RuntimeError(
-            "all values given when constructing a transformation matrix must have the same type"
-        ) from None
-
-    if dtype is None:
-        dtype = value_type
-    elif not warp.types.scalars_equal(value_type, dtype):
-        raise RuntimeError(
-            f"all values used to initialize this transformation matrix are expected to be of the type `{dtype.__name__}`"
-        )
-
-    return matrix(shape=(4, 4), dtype=dtype)
-
-
-def matrix_transform_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
-    # We're in the codegen stage where we emit the code calling the built-in.
-    # Further validate the given argument values if needed and map them
-    # to the underlying C++ function's runtime and template params.
-
-    dtype = return_type._wp_scalar_type_
-
-    func_args = tuple(v for k, v in args.items() if k != "dtype")
-    template_args = (4, 4, dtype)
-    return (func_args, template_args)
+    raise RuntimeError(
+        "the built-in `wp.matrix()` to construct a 4x4 matrix from a 3D position, quaternion, "
+        "and 3D scale vector has been removed in favor of `wp.transform_compose()`."
+    )
 
 
 add_builtin(
@@ -1259,13 +1259,14 @@ add_builtin(
     defaults={"dtype": None},
     value_func=matrix_transform_value_func,
     export_func=lambda input_types: {k: v for k, v in input_types.items() if k != "dtype"},
-    dispatch_func=matrix_transform_dispatch_func,
     native_func="mat_t",
     doc="""Construct a 4x4 transformation matrix that applies the transformations as
     Translation(pos)*Rotation(rot)*Scaling(scale) when applied to column vectors, i.e.: y = (TRS)*x
 
-    .. warning::
-       This function has been deprecated in favor of :func:`warp.math.transform_compose()`.""",
+    .. versionremoved:: 1.10
+       This function has been removed in favor of :func:`warp.math.transform_compose()`.
+
+    .. deprecated:: 1.8""",
     group="Vector Math",
     export=False,
 )
@@ -1460,7 +1461,7 @@ def quaternion_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str
 
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"all values used to initialize this quaternion are expected to be of the type `{dtype.__name__}`"
             )
@@ -1568,7 +1569,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -1697,7 +1698,7 @@ def transformation_value_func(arg_types: Mapping[str, type], arg_values: Mapping
         value_type = strip_reference(variadic_arg_types[0])
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"the value used to fill this transform is expected to be of the type `{dtype.__name__}`"
             )
@@ -1710,7 +1711,7 @@ def transformation_value_func(arg_types: Mapping[str, type], arg_values: Mapping
 
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"all values used to initialize this transform are expected to be of the type `{dtype.__name__}`"
             )
@@ -1735,7 +1736,7 @@ def transformation_pq_value_func(arg_types: Mapping[str, type], arg_values: Mapp
     dtype = arg_values.get("dtype", None)
     if dtype is None:
         dtype = value_type
-    elif not warp.types.scalars_equal(value_type, dtype):
+    elif not warp._src.types.scalars_equal(value_type, dtype):
         raise RuntimeError(
             f"all values used to initialize this transformation matrix are expected to be of the type `{dtype.__name__}`"
         )
@@ -1750,9 +1751,19 @@ def transformation_dispatch_func(input_types: Mapping[str, type], return_type: A
 
     dtype = return_type._wp_scalar_type_
 
-    variadic_args = tuple(v for k, v in args.items() if k != "dtype")
+    variadic_args = args.get("args", ())
+    variadic_arg_count = len(variadic_args)
+
+    if variadic_arg_count == 7:
+        func_args = variadic_args
+    else:
+        func_args = tuple(v for k, v in args.items() if k != "dtype")
+        if "p" in args and "q" not in args:
+            quat_ident = warp._src.codegen.Var(
+                label=None, type=quaternion(dtype=dtype), constant=quaternion(dtype=dtype)(0, 0, 0, 1)
+            )
+            func_args += (quat_ident,)
 
-    func_args = variadic_args
     template_args = (dtype,)
     return (func_args, template_args)
 
@@ -1760,7 +1771,7 @@ def transformation_dispatch_func(input_types: Mapping[str, type], return_type: A
 add_builtin(
     "transformation",
     input_types={"p": vector(length=3, dtype=Float), "q": quaternion(dtype=Float), "dtype": Float},
-    defaults={"dtype": None},
+    defaults={"q": None, "dtype": None},
     value_func=transformation_pq_value_func,
     export_func=lambda input_types: {k: v for k, v in input_types.items() if k != "dtype"},
     dispatch_func=transformation_dispatch_func,
@@ -1784,7 +1795,6 @@ add_builtin(
     doc="Construct a spatial transform vector of given dtype.",
     group="Spatial Math",
     export=False,
-    missing_grad=True,
 )
 
 
@@ -1819,7 +1829,7 @@ add_builtin(
     group="Transformations",
     doc="Construct an identity transform with zero translation and identity rotation.",
     export=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -1953,7 +1963,7 @@ def spatial_vector_value_func(arg_types: Mapping[str, type], arg_values: Mapping
 
         if dtype is None:
             dtype = value_type
-        elif not warp.types.scalars_equal(value_type, dtype):
+        elif not warp._src.types.scalars_equal(value_type, dtype):
             raise RuntimeError(
                 f"all values used to initialize this spatial vector are expected to be of the type `{dtype.__name__}`"
             )
@@ -2147,7 +2157,7 @@ add_builtin(
     value_func=tile_zeros_value_func,
     dispatch_func=tile_zeros_dispatch_func,
     variadic=False,
-    missing_grad=True,
+    is_differentiable=False,
     doc="""Allocate a tile of zero-initialized items.
 
     :param shape: Shape of the output tile
@@ -2167,7 +2177,7 @@ add_builtin(
     value_func=tile_zeros_value_func,
     dispatch_func=tile_zeros_dispatch_func,
     variadic=False,
-    missing_grad=True,
+    is_differentiable=False,
     hidden=True,
     group="Tile Primitives",
     export=False,
@@ -2219,7 +2229,7 @@ add_builtin(
     defaults={"storage": "register"},
     value_func=tile_ones_value_func,
     dispatch_func=tile_ones_dispatch_func,
-    missing_grad=True,
+    is_differentiable=False,
     doc="""Allocate a tile of one-initialized items.
 
     :param shape: Shape of the output tile
@@ -2238,7 +2248,86 @@ add_builtin(
     defaults={"storage": "register"},
     value_func=tile_ones_value_func,
     dispatch_func=tile_ones_dispatch_func,
-    missing_grad=True,
+    is_differentiable=False,
+    hidden=True,
+    group="Tile Primitives",
+    export=False,
+)
+
+
+def tile_full_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
+    # return generic type (for doc builds)
+    if arg_types is None:
+        return tile(dtype=Any, shape=Tuple[int, ...])
+
+    shape = extract_tuple(arg_values["shape"], as_constant=True)
+
+    if None in shape:
+        raise ValueError("Tile functions require shape to be a compile time constant.")
+
+    if "value" not in arg_values:
+        raise TypeError("tile_full() missing required keyword argument 'value'")
+
+    if "dtype" not in arg_values:
+        raise TypeError("tile_full() missing required keyword argument 'dtype'")
+
+    if "storage" not in arg_values:
+        raise TypeError("tile_full() missing required keyword argument 'storage'")
+
+    if arg_values["storage"] not in {"shared", "register"}:
+        raise ValueError(f"Invalid value for 'storage': {arg_values['storage']!r}. Expected 'shared' or 'register'.")
+
+    dtype = arg_values["dtype"]
+
+    return tile(dtype=dtype, shape=shape, storage=arg_values["storage"])
+
+
+def tile_full_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
+    shape = extract_tuple(arg_values["shape"], as_constant=True)
+
+    if None in shape:
+        raise ValueError("Tile functions require shape to be a compile time constant.")
+
+    dtype = arg_values["dtype"]
+    value = arg_values["value"]
+
+    func_args = [value]
+
+    template_args = []
+    template_args.append(dtype)
+    template_args.extend(shape)
+
+    return (func_args, template_args)
+
+
+add_builtin(
+    "tile_full",
+    input_types={"shape": Tuple[int, ...], "value": Any, "dtype": Any, "storage": str},
+    defaults={"storage": "register"},
+    value_func=tile_full_value_func,
+    dispatch_func=tile_full_dispatch_func,
+    is_differentiable=False,
+    doc="""Allocate a tile filled with the specified value.
+
+    :param shape: Shape of the output tile
+    :param value: Value to fill the tile with
+    :param dtype: Data type of output tile's elements
+    :param storage: The storage location for the tile: ``"register"`` for registers
+      (default) or ``"shared"`` for shared memory.
+    :returns: A tile filled with the specified value""",
+    group="Tile Primitives",
+    export=False,
+)
+
+
+# overload for scalar shape
+add_builtin(
+    "tile_full",
+    input_types={"shape": int, "value": Any, "dtype": Any, "storage": str},
+    defaults={"storage": "register"},
+    value_func=tile_full_value_func,
+    dispatch_func=tile_full_dispatch_func,
+    is_differentiable=False,
     hidden=True,
     group="Tile Primitives",
     export=False,
@@ -2300,13 +2389,13 @@ def tile_arange_dispatch_func(arg_types: Mapping[str, type], return_type: Any, a
     args = arg_values["args"]
 
     if len(args) == 1:
-        start = warp.codegen.Var(label=None, type=return_type.dtype, constant=0)
+        start = warp._src.codegen.Var(label=None, type=return_type.dtype, constant=0)
         stop = args[0]
-        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=1)
+        step = warp._src.codegen.Var(label=None, type=return_type.dtype, constant=1)
     elif len(args) == 2:
         start = args[0]
         stop = args[1]
-        step = warp.codegen.Var(label=None, type=return_type.dtype, constant=1)
+        step = warp._src.codegen.Var(label=None, type=return_type.dtype, constant=1)
     elif len(args) == 3:
         start = args[0]
         stop = args[1]
@@ -2329,7 +2418,7 @@ add_builtin(
     value_func=tile_arange_value_func,
     dispatch_func=tile_arange_dispatch_func,
     variadic=True,
-    missing_grad=True,
+    is_differentiable=False,
     doc="""Generate a tile of linearly spaced elements.
 
     :param args: Variable-length positional arguments, interpreted as:
@@ -3124,7 +3213,7 @@ add_builtin(
     :param shape: Shape of the returned slice
     :returns: A tile with dimensions given by the specified shape or the remaining source tile dimensions""",
     group="Tile Primitives",
-    missing_grad=True,
+    is_differentiable=False,
     export=False,
 )
 
@@ -3371,7 +3460,32 @@ add_builtin(
 
 add_builtin(
     "assign",
-    input_types={"dst": tile(dtype=Any, shape=Tuple[int, int]), "i": int, "j": int, "src": Any},
+    input_types={"dst": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "src": Any},
+    value_func=tile_assign_value_func,
+    group="Tile Primitives",
+    export=False,
+    hidden=True,
+)
+
+add_builtin(
+    "assign",
+    input_types={"dst": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "src": Any},
+    value_func=tile_assign_value_func,
+    group="Tile Primitives",
+    export=False,
+    hidden=True,
+)
+
+add_builtin(
+    "assign",
+    input_types={
+        "dst": tile(dtype=Any, shape=Tuple[int, ...]),
+        "i": int,
+        "j": int,
+        "k": int,
+        "l": int,
+        "src": Any,
+    },
     value_func=tile_assign_value_func,
     group="Tile Primitives",
     export=False,
@@ -3380,7 +3494,15 @@ add_builtin(
 
 add_builtin(
     "assign",
-    input_types={"dst": tile(dtype=Any, shape=Tuple[int, int, int]), "i": int, "j": int, "k": int, "src": Any},
+    input_types={
+        "dst": tile(dtype=Any, shape=Tuple[int, ...]),
+        "i": int,
+        "j": int,
+        "k": int,
+        "l": int,
+        "m": int,
+        "src": Any,
+    },
     value_func=tile_assign_value_func,
     group="Tile Primitives",
     export=False,
@@ -3395,6 +3517,8 @@ add_builtin(
         "j": int,
         "k": int,
         "l": int,
+        "m": int,
+        "n": int,
         "src": Any,
     },
     value_func=tile_assign_value_func,
@@ -3416,7 +3540,7 @@ def tile_value_func(arg_types, arg_values):
 
     if preserve_type:
         dtype = arg_types["x"]
-        shape = (warp.codegen.options["block_dim"],)
+        shape = (warp._src.codegen.options["block_dim"],)
 
         return tile(dtype=dtype, shape=shape)
 
@@ -3424,18 +3548,18 @@ def tile_value_func(arg_types, arg_values):
         if type_is_vector(arg_types["x"]):
             dtype = arg_types["x"]._wp_scalar_type_
             length = arg_types["x"]._shape_[0]
-            shape = (length, warp.codegen.options["block_dim"])
+            shape = (length, warp._src.codegen.options["block_dim"])
         elif type_is_quaternion(arg_types["x"]):
             dtype = arg_types["x"]._wp_scalar_type_
-            shape = (4, warp.codegen.options["block_dim"])
+            shape = (4, warp._src.codegen.options["block_dim"])
         elif type_is_matrix(arg_types["x"]):
             dtype = arg_types["x"]._wp_scalar_type_
             rows = arg_types["x"]._shape_[0]
             cols = arg_types["x"]._shape_[1]
-            shape = (rows, cols, warp.codegen.options["block_dim"])
+            shape = (rows, cols, warp._src.codegen.options["block_dim"])
         else:
             dtype = arg_types["x"]
-            shape = (warp.codegen.options["block_dim"],)
+            shape = (warp._src.codegen.options["block_dim"],)
 
         return tile(dtype=dtype, shape=shape)
 
@@ -3525,17 +3649,17 @@ def untile_value_func(arg_types, arg_values):
     if not is_tile(t):
         raise TypeError(f"untile() argument must be a tile, got {t!r}")
 
-    if t.shape[-1] != warp.codegen.options["block_dim"]:
+    if t.shape[-1] != warp._src.codegen.options["block_dim"]:
         raise ValueError(
-            f"untile() argument last dimension {t.shape[-1]} does not match the expected block width {warp.codegen.options['block_dim']}"
+            f"untile() argument last dimension {t.shape[-1]} does not match the expected block width {warp._src.codegen.options['block_dim']}"
         )
 
     if len(t.shape) == 1:
         return t.dtype
     elif len(t.shape) == 2:
-        return warp.types.vector(t.shape[0], t.dtype)
+        return warp._src.types.vector(t.shape[0], t.dtype)
     elif len(t.shape) == 3:
-        return warp.types.matrix((t.shape[0], t.shape[1]), t.dtype)
+        return warp._src.types.matrix((t.shape[0], t.shape[1]), t.dtype)
     else:
         raise ValueError(f"untile() argument must have a positive size in dimension 0, but got {t.shape[0]}")
 
@@ -3597,7 +3721,36 @@ def tile_extract_value_func(arg_types, arg_values):
     # force the input tile to shared memory
     arg_types["a"].storage = "shared"
 
-    return arg_types["a"].dtype
+    # count the number of indices (all parameters except the tile "a")
+    num_indices = len(arg_types) - 1
+    tile_dtype = arg_types["a"].dtype
+    tile_shape = arg_types["a"].shape
+
+    if type_is_vector(tile_dtype):
+        if num_indices == len(tile_shape):
+            return tile_dtype
+        elif num_indices == len(tile_shape) + 1:
+            return tile_dtype._wp_scalar_type_
+        else:
+            raise IndexError(
+                f"tile_extract: incorrect number of indices ({num_indices}) for tile shape {tuple(tile_shape)}"
+            )
+    elif type_is_matrix(tile_dtype):
+        if num_indices == len(tile_shape):
+            return tile_dtype
+        elif num_indices == len(tile_shape) + 2:
+            return tile_dtype._wp_scalar_type_
+        else:
+            raise IndexError(
+                f"tile_extract: incorrect number of indices ({num_indices}) for matrix tile shape {tuple(tile_shape)}"
+            )
+    else:
+        # scalar element: index count must exactly match tile rank
+        if num_indices == len(tile_shape):
+            return tile_dtype
+        raise IndexError(
+            f"tile_extract: incorrect number of indices ({num_indices}) for tile shape {tuple(tile_shape)}"
+        )
 
 
 add_builtin(
@@ -3621,7 +3774,7 @@ add_builtin(
 
 add_builtin(
     "tile_extract",
-    input_types={"a": tile(dtype=Any, shape=Tuple[int, int]), "i": int, "j": int},
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int},
     value_func=tile_extract_value_func,
     variadic=False,
     doc="""Extract a single element from the tile.
@@ -3632,7 +3785,7 @@ add_builtin(
 
     :param a: Tile to extract the element from
     :param i: Coordinate of element on first dimension
-    :param j: Coordinate of element on the second dimension
+    :param j: Coordinate of element on the second dimension, or vector index
     :returns: The value of the element at the specified tile location with the same data type as the input tile""",
     group="Tile Primitives",
     hidden=True,
@@ -3641,7 +3794,57 @@ add_builtin(
 
 add_builtin(
     "tile_extract",
-    input_types={"a": tile(dtype=Any, shape=Tuple[int, int, int]), "i": int, "j": int, "k": int},
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int},
+    value_func=tile_extract_value_func,
+    variadic=False,
+    doc="""Extract a single element from the tile.
+
+    This function will extract an element from the tile and broadcast its value to all threads in the block.
+
+    Note that this may incur additional synchronization if the source tile is a register tile.
+
+    :param a: Tile to extract the element from
+    :param i: Coordinate of element on first dimension
+    :param j: Coordinate of element on the second dimension, or first matrix index
+    :param k: Coordinate of element on the third dimension, or vector index, or second matrix index
+    :returns: The value of the element at the specified tile location with the same data type as the input tile""",
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
+add_builtin(
+    "tile_extract",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "l": int},
+    value_func=tile_extract_value_func,
+    variadic=False,
+    doc="""Extract a single element from the tile.
+
+    This function will extract an element from the tile and broadcast its value to all threads in the block.
+
+    Note that this may incur additional synchronization if the source tile is a register tile.
+
+    :param a: Tile to extract the element from
+    :param i: Coordinate of element on first dimension
+    :param j: Coordinate of element on the second dimension
+    :param k: Coordinate of element on the third dimension, or first matrix index
+    :param l: Coordinate of element on the fourth dimension, or vector index, or second matrix index
+    :returns: The value of the element at the specified tile location, with the same data type as the input tile""",
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+)
+
+add_builtin(
+    "tile_extract",
+    input_types={
+        "a": tile(dtype=Any, shape=Tuple[int, ...]),
+        "i": int,
+        "j": int,
+        "k": int,
+        "l": int,
+        "m": int,
+    },
     value_func=tile_extract_value_func,
     variadic=False,
     doc="""Extract a single element from the tile.
@@ -3654,7 +3857,9 @@ add_builtin(
     :param i: Coordinate of element on first dimension
     :param j: Coordinate of element on the second dimension
     :param k: Coordinate of element on the third dimension
-    :returns: The value of the element at the specified tile location with the same data type as the input tile""",
+    :param l: Coordinate of element on the fourth dimension, or first matrix index
+    :param m: Vector index, or second matrix index
+    :returns: The value of the element at the specified tile location, with the same data type as the input tile""",
     group="Tile Primitives",
     hidden=True,
     export=False,
@@ -3662,7 +3867,15 @@ add_builtin(
 
 add_builtin(
     "tile_extract",
-    input_types={"a": tile(dtype=Any, shape=Tuple[int, int, int, int]), "i": int, "j": int, "k": int, "l": int},
+    input_types={
+        "a": tile(dtype=Any, shape=Tuple[int, int, int, int]),
+        "i": int,
+        "j": int,
+        "k": int,
+        "l": int,
+        "m": int,
+        "n": int,
+    },
     value_func=tile_extract_value_func,
     variadic=False,
     doc="""Extract a single element from the tile.
@@ -3676,6 +3889,8 @@ add_builtin(
     :param j: Coordinate of element on the second dimension
     :param k: Coordinate of element on the third dimension
     :param l: Coordinate of element on the fourth dimension
+    :param m: Vector index, or first matrix index
+    :param n: Second matrix index
     :returns: The value of the element at the specified tile location, with the same data type as the input tile""",
     group="Tile Primitives",
     hidden=True,
@@ -3762,49 +3977,160 @@ add_builtin(
     export=False,
 )
 
-
-def tile_transpose_value_func(arg_types, arg_values):
-    # return generic type (for doc builds)
-    if arg_types is None:
-        return tile(dtype=Any, shape=Tuple[int, int])
-
-    if len(arg_types) != 1:
-        raise TypeError(f"tile_transpose() takes exactly 1 positional argument but {len(arg_types)} were given")
-
-    t = arg_types["a"]
-
-    if not is_tile(t):
-        raise TypeError(f"tile_transpose() argument must be a tile, got {t!r}")
-
-    layout = None
-
-    # flip layout
-    if t.layout == "rowmajor":
-        layout = "colmajor"
-    elif t.layout == "colmajor":
-        layout = "rowmajor"
-
-    # force the input tile to shared memory
-    t.storage = "shared"
-
-    return tile(
-        dtype=t.dtype,
-        shape=t.shape[::-1],
-        storage=t.storage,
-        strides=t.strides[::-1],
-        layout=layout,
-        owner=False,
-    )
-
-
 add_builtin(
-    "tile_transpose",
-    input_types={"a": tile(dtype=Any, shape=Tuple[int, int])},
-    value_func=tile_transpose_value_func,
-    variadic=True,
-    doc="""Transpose a tile.
-
-    For shared memory tiles, this operation will alias the input tile.
+    "tile_bit_and_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_and_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_and_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_and_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "l": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+
+add_builtin(
+    "tile_bit_or_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_or_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_or_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_or_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "l": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+
+add_builtin(
+    "tile_bit_xor_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_xor_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_xor_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+add_builtin(
+    "tile_bit_xor_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "i": int, "j": int, "k": int, "l": int, "value": Any},
+    value_func=tile_inplace_value_func,
+    group="Tile Primitives",
+    hidden=True,
+    export=False,
+    is_differentiable=False,
+)
+
+
+def tile_transpose_value_func(arg_types, arg_values):
+    # return generic type (for doc builds)
+    if arg_types is None:
+        return tile(dtype=Any, shape=Tuple[int, int])
+
+    if len(arg_types) != 1:
+        raise TypeError(f"tile_transpose() takes exactly 1 positional argument but {len(arg_types)} were given")
+
+    t = arg_types["a"]
+
+    if not is_tile(t):
+        raise TypeError(f"tile_transpose() argument must be a tile, got {t!r}")
+
+    layout = None
+
+    # flip layout
+    if t.layout == "rowmajor":
+        layout = "colmajor"
+    elif t.layout == "colmajor":
+        layout = "rowmajor"
+
+    # force the input tile to shared memory
+    t.storage = "shared"
+
+    return tile(
+        dtype=t.dtype,
+        shape=t.shape[::-1],
+        storage=t.storage,
+        strides=t.strides[::-1],
+        layout=layout,
+        owner=False,
+    )
+
+
+add_builtin(
+    "tile_transpose",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, int])},
+    value_func=tile_transpose_value_func,
+    variadic=True,
+    doc="""Transpose a tile.
+
+    For shared memory tiles, this operation will alias the input tile.
     Register tiles will first be transferred to shared memory before transposition.
 
     :param a: Tile to transpose with ``shape=(M,N)``
@@ -3935,6 +4261,80 @@ add_builtin(
 )
 
 
+def tile_sum_axis_value_func(arg_types, arg_values):
+    if arg_types is None:
+        return tile(dtype=Scalar, shape=Tuple[int, ...])
+
+    a = arg_types["a"]
+
+    if not is_tile(a):
+        raise TypeError(f"tile_sum() 'a' argument must be a tile, got {a!r}")
+
+    # force input tile to shared
+    a.storage = "shared"
+
+    axis = arg_values["axis"]
+    shape = a.shape
+
+    if axis < 0 or axis >= len(shape):
+        raise ValueError(f"tile_sum() axis {axis} is out of bounds for tile with {len(shape)} dimensions")
+
+    # shape is identical less the axis reduction is along
+    if len(shape) > 1:
+        new_shape = shape[:axis] + shape[axis + 1 :]
+    else:
+        new_shape = (1,)
+
+    return tile(dtype=a.dtype, shape=new_shape)
+
+
+def tile_sum_axis_dispatch_func(arg_types: Mapping[str, type], return_type: Any, arg_values: Mapping[str, Var]):
+    tile = arg_values["a"]
+    axis_var = arg_values["axis"]
+    if not hasattr(axis_var, "constant") or axis_var.constant is None:
+        raise ValueError("tile_sum() axis must be a compile-time constant")
+    axis = axis_var.constant
+
+    return ((tile,), (axis,))
+
+
+add_builtin(
+    "tile_sum",
+    input_types={"a": tile(dtype=Scalar, shape=Tuple[int, ...]), "axis": int},
+    value_func=tile_sum_axis_value_func,
+    dispatch_func=tile_sum_axis_dispatch_func,
+    doc="""Cooperatively compute the sum of the tile elements across an axis of the tile using all threads in the block.
+
+    :param a: The input tile. Must reside in shared memory.
+    :param axis: The tile axis to compute the sum across. Must be a compile-time constant.
+    :returns: A tile with the same shape as the input tile less the axis dimension and the same data type as the input tile.
+
+    Example:
+
+    .. code-block:: python
+
+        @wp.kernel
+        def compute():
+
+            t = wp.tile_ones(dtype=float, shape=(8, 8))
+            s = wp.tile_sum(t, axis=0)
+
+            print(s)
+
+        wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
+
+    Prints:
+
+    .. code-block:: text
+
+        [8 8 8 8 8 8 8 8] = tile(shape=(8), storage=register)
+
+    """,
+    group="Tile Primitives",
+    export=False,
+)
+
+
 def tile_sort_value_func(arg_types, arg_values):
     # return generic type (for doc builds)
     if arg_types is None:
@@ -4011,7 +4411,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4065,7 +4465,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4119,7 +4519,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4172,7 +4572,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4225,11 +4625,10 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
-# does type propagation for load()
 def tile_reduce_value_func(arg_types, arg_values):
     if arg_types is None:
         return tile(dtype=Scalar, shape=(1,))
@@ -4283,7 +4682,88 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
+)
+
+
+def tile_reduce_axis_value_func(arg_types, arg_values):
+    if arg_types is None:
+        return tile(dtype=Scalar, shape=Tuple[int, ...])
+
+    a = arg_types["a"]
+
+    if not is_tile(a):
+        raise TypeError(f"tile_reduce() 'a' argument must be a tile, got {a!r}")
+
+    # force input tile to shared memory
+    a.storage = "shared"
+
+    axis = arg_values["axis"]
+    shape = a.shape
+
+    if axis < 0 or axis >= len(shape):
+        raise ValueError(f"tile_reduce() axis {axis} is out of bounds for tile with {len(shape)} dimensions")
+
+    # shape is identical less the axis reduction is along
+    if len(shape) > 1:
+        new_shape = shape[:axis] + shape[axis + 1 :]
+    else:
+        new_shape = (1,)
+
+    return tile(dtype=a.dtype, shape=new_shape)
+
+
+add_builtin(
+    "tile_reduce",
+    input_types={"op": Callable, "a": tile(dtype=Scalar, shape=Tuple[int, ...]), "axis": int},
+    value_func=tile_reduce_axis_value_func,
+    native_func="tile_reduce_axis",
+    doc="""Apply a custom reduction operator across a tile axis.
+
+    This function cooperatively performs a reduction using the provided operator across an axis of the tile.
+
+    :param op: A callable function that accepts two arguments and returns one argument, may be a user function or builtin
+    :param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type. Must reside in shared memory.
+    :param axis: The tile axis to perform the reduction across. Must be a compile-time constant.
+    :returns: A tile with the same shape as the input tile less the axis dimension and the same data type as the input tile.
+
+    Example:
+
+    .. code-block:: python
+
+        TILE_M = wp.constant(4)
+        TILE_N = wp.constant(2)
+
+        @wp.kernel
+        def compute(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+
+            a = wp.tile_load(x, shape=(TILE_M, TILE_N))
+            b = wp.tile_reduce(wp.add, a, axis=1)
+            wp.tile_store(y, b)
+
+        arr = np.arange(TILE_M * TILE_N).reshape(TILE_M, TILE_N)
+
+        x = wp.array(arr, dtype=float)
+        y = wp.zeros(TILE_M, dtype=float)
+
+        wp.launch_tiled(compute, dim=[1], inputs=[x], outputs=[y], block_dim=32)
+
+        print(x.numpy())
+        print(y.numpy())
+
+    Prints:
+
+    .. code-block:: text
+
+        [[0. 1.]
+         [2. 3.]
+         [4. 5.]
+         [6. 7.]]
+        [ 1.  5.  9. 13.]
+    """,
+    group="Tile Primitives",
+    export=False,
+    is_differentiable=False,
 )
 
 
@@ -4347,7 +4827,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4411,7 +4891,7 @@ add_builtin(
     """,
     group="Tile Primitives",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -4665,7 +5145,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -4681,7 +5161,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -4691,7 +5171,7 @@ add_builtin(
     doc="WIP",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -4743,7 +5223,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -4759,7 +5239,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -4770,7 +5250,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5111,7 +5591,7 @@ add_builtin(
     :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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5123,7 +5603,7 @@ add_builtin(
 
     The index of the current face is stored in ``index``, returns ``False`` if there are no more overlapping triangles.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5153,7 +5633,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5165,7 +5645,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5179,7 +5659,7 @@ add_builtin(
 
     Returns -1 if the :class:`HashGrid` has not been reserved.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5189,16 +5669,34 @@ add_builtin(
     group="Geometry",
     doc="""Tests for intersection between two triangles (v0, v1, v2) and (u0, u1, u2) using Moller's method.
 
+    This function works with single precision, may return incorrect results in some case.
+
+    Returns > 0 if triangles intersect.""",
+    export=False,
+    is_differentiable=False,
+)
+
+
+add_builtin(
+    "intersect_tri_tri",
+    input_types={"v0": vec3d, "v1": vec3d, "v2": vec3d, "u0": vec3d, "u1": vec3d, "u2": vec3d},
+    value_type=int,
+    group="Geometry",
+    doc="""Tests for intersection between two triangles (v0, v1, v2) and (u0, u1, u2) using Moller's method.
+
+    This function works with double precision, results are more accurate than the single precision version.
+
     Returns > 0 if triangles intersect.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
+
 add_builtin(
     "mesh_get",
     input_types={"id": uint64},
     value_type=Mesh,
-    missing_grad=True,
+    is_differentiable=False,
     group="Geometry",
     doc="""Retrieves the mesh given its index.""",
     export=False,
@@ -5211,7 +5709,7 @@ add_builtin(
     group="Geometry",
     doc="""Evaluates the face normal the mesh given a face index.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5221,7 +5719,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the point of the mesh given a index.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5231,7 +5729,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the velocity of the mesh given a index.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5241,7 +5739,7 @@ add_builtin(
     group="Geometry",
     doc="""Returns the point-index of the mesh given a face-vertex index.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -5289,7 +5787,7 @@ add_builtin(
     group="Utility",
     export=False,
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "iter_next",
@@ -5298,7 +5796,7 @@ add_builtin(
     group="Utility",
     export=False,
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "iter_next",
@@ -5307,7 +5805,7 @@ add_builtin(
     group="Utility",
     export=False,
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5318,7 +5816,7 @@ add_builtin(
     group="Utility",
     doc="""Returns the range in reversed order.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 # ---------------------------------
@@ -5338,8 +5836,8 @@ _volume_supported_value_types = {
 
 
 def _is_volume_type_supported(dtype):
-    for typ in _volume_supported_value_types:
-        if types_equal(typ, dtype):
+    for value_type in _volume_supported_value_types:
+        if types_equal(value_type, dtype):
             return True
     return False
 
@@ -5467,7 +5965,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,
+    is_differentiable=False,
 )
 
 
@@ -5488,7 +5986,7 @@ add_builtin(
     export=False,
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5519,7 +6017,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5528,7 +6026,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5549,7 +6047,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5558,7 +6056,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5577,7 +6075,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5586,7 +6084,7 @@ add_builtin(
     group="Volumes",
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -5668,7 +6166,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5710,7 +6208,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5722,7 +6220,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5731,7 +6229,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Return a random integer in the range [-2^31, 2^31).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randi",
@@ -5739,7 +6237,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Return a random integer between [low, high).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randu",
@@ -5747,7 +6245,7 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="Return a random unsigned integer in the range [0, 2^32).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randu",
@@ -5755,7 +6253,7 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="Return a random unsigned integer between [low, high).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randf",
@@ -5763,7 +6261,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Return a random float between [0.0, 1.0).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randf",
@@ -5771,7 +6269,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Return a random float between [low, high).",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "randn",
@@ -5779,7 +6277,7 @@ add_builtin(
     value_type=float,
     group="Random",
     doc="Sample a normal (Gaussian) distribution of mean 0 and variance 1. ",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5788,7 +6286,7 @@ add_builtin(
     value_type=int,
     group="Random",
     doc="Inverse-transform sample a cumulative distribution function.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_triangle",
@@ -5796,7 +6294,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a triangle. Returns sample barycentric coordinates.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_ring",
@@ -5804,7 +6302,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a ring in the xy plane.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_disk",
@@ -5812,7 +6310,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a disk in the xy plane.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_sphere_surface",
@@ -5820,7 +6318,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit sphere surface.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_sphere",
@@ -5828,7 +6326,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit sphere.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_hemisphere_surface",
@@ -5836,7 +6334,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit hemisphere surface.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_hemisphere",
@@ -5844,7 +6342,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit hemisphere.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_square",
@@ -5852,7 +6350,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Uniformly sample a unit square.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "sample_unit_cube",
@@ -5860,7 +6358,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Uniformly sample a unit cube.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5872,7 +6370,7 @@ add_builtin(
 
     :param state: RNG state
     :param lam: The expected value of the distribution""",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -5940,7 +6438,7 @@ add_builtin(
     value_type=vec2,
     group="Random",
     doc="Divergence-free vector field based on the gradient of a Perlin noise function.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "curlnoise",
@@ -5949,7 +6447,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Divergence-free vector field based on the curl of three Perlin noise functions.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "curlnoise",
@@ -5958,7 +6456,7 @@ add_builtin(
     value_type=vec3,
     group="Random",
     doc="Divergence-free vector field based on the curl of three Perlin noise functions.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -5990,7 +6488,7 @@ add_builtin(
     dispatch_func=printf_dispatch_func,
     group="Utility",
     doc="Allows printing formatted strings using C-style format specifiers.",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -6009,7 +6507,7 @@ add_builtin(
     group="Utility",
     namespace="",
     native_func="__debugbreak",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 # helpers
@@ -6027,7 +6525,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid1d",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -6038,7 +6536,7 @@ add_builtin(
     doc="Returns the number of threads in the current block.",
     namespace="",
     native_func="builtin_block_dim",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -6053,7 +6551,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid2d",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -6068,7 +6566,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid3d",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -6083,7 +6581,7 @@ add_builtin(
     This function may not be called from user-defined Warp functions.""",
     namespace="",
     native_func="builtin_tid4d",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -6127,56 +6625,20 @@ def select_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, An
     if arg_types is None:
         return Any
 
-    v_true = arg_types["value_if_true"]
-    v_false = arg_types["value_if_false"]
+    raise RuntimeError("wp.select() has been removed. Use wp.where(cond, value_if_true, value_if_false) instead.")
 
-    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
+add_builtin(
+    "select",
+    input_types={"cond": builtins.bool, "value_if_false": Any, "value_if_true": Any},
+    value_func=select_value_func,
+    doc="""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
 
-        # 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"
-        "version. Use wp.where(cond, value_if_true, value_if_false) instead.",
-        category=DeprecationWarning,
-    )
-
-    func_args = tuple(args.values())
-    template_args = ()
-
-    return (func_args, template_args)
-
-
-add_builtin(
-    "select",
-    input_types={"cond": builtins.bool, "value_if_false": Any, "value_if_true": Any},
-    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``.
-
-    .. deprecated:: 1.7
+    .. versionremoved:: 1.10
          Use :func:`where` instead, which has the more intuitive argument order:
-         ``where(cond, value_if_true, value_if_false)``.""",
+         ``where(cond, value_if_true, value_if_false)``.
+
+    .. deprecated:: 1.7""",
     group="Utility",
 )
 for t in int_types:
@@ -6184,24 +6646,26 @@ for t in int_types:
         "select",
         input_types={"cond": t, "value_if_false": Any, "value_if_true": Any},
         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``.
 
-    .. deprecated:: 1.7
+    .. versionremoved:: 1.10
          Use :func:`where` instead, which has the more intuitive argument order:
-         ``where(cond, value_if_true, value_if_false)``.""",
+         ``where(cond, value_if_true, value_if_false)``.
+
+    .. deprecated:: 1.7""",
         group="Utility",
     )
 add_builtin(
     "select",
     input_types={"arr": array(dtype=Any), "value_if_false": Any, "value_if_true": Any},
     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``.
 
-    .. deprecated:: 1.7
+    .. versionremoved:: 1.10
          Use :func:`where` instead, which has the more intuitive argument order:
-         ``where(arr, value_if_true, value_if_false)``.""",
+         ``where(arr, value_if_true, value_if_false)``.
+
+    .. deprecated:: 1.7""",
     group="Utility",
 )
 
@@ -6291,7 +6755,7 @@ add_builtin(
     group="Utility",
     hidden=True,
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -6332,7 +6796,7 @@ add_builtin(
     native_func="fixedarray_t",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
     hidden=True,  # Unhide once we can document both a built-in and a Python scope function sharing the same name.
 )
 
@@ -6375,14 +6839,13 @@ for array_type in array_types:
 # does argument checking and type propagation for view()
 def view_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     arr_type = arg_types["arr"]
-    idx_types = tuple(arg_types[x] for x in "ijk" if arg_types.get(x, None) is not None)
+    idx_types = tuple(arg_types[x] for x in "ijkl" if arg_types.get(x, None) is not None)
 
     if not is_array(arr_type):
         raise RuntimeError("view() first argument must be an array")
 
     idx_count = len(idx_types)
-
-    if idx_count >= arr_type.ndim:
+    if idx_count > arr_type.ndim:
         raise RuntimeError(
             f"Trying to create an array view with {idx_count} indices, "
             f"but the array only has {arr_type.ndim} dimension(s). "
@@ -6390,14 +6853,35 @@ def view_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]
             f"the expected number of dimensions, e.g.: def func(param: wp.array3d(dtype=float): ..."
         )
 
-    # check index types
-    for t in idx_types:
-        if not type_is_int(t):
-            raise RuntimeError(f"view() index arguments must be of integer type, got index of type {type_repr(t)}")
+    has_slice = any(is_slice(x) for x in idx_types)
+    if has_slice:
+        # check index types
+        for t in idx_types:
+            if not (type_is_int(t) or is_slice(t)):
+                raise RuntimeError(
+                    f"view() index arguments must be of integer or slice types, got index of type {type_repr(t)}"
+                )
+
+        # Each integer index collapses one dimension.
+        int_count = sum(x.step == 0 for x in idx_types)
+        ndim = arr_type.ndim - int_count
+        assert ndim > 0
+    else:
+        if idx_count == arr_type.ndim:
+            raise RuntimeError("Expected to call `address()` instead of `view()`")
+
+        # check index types
+        for t in idx_types:
+            if not type_is_int(t):
+                raise RuntimeError(
+                    f"view() index arguments must be of integer or slice types, got index of type {type_repr(t)}"
+                )
+
+        # create an array view with leading dimensions removed
+        ndim = arr_type.ndim - idx_count
+        assert ndim > 0
 
-    # create an array view with leading dimensions removed
     dtype = arr_type.dtype
-    ndim = arr_type.ndim - idx_count
     if isinstance(arr_type, (fabricarray, indexedfabricarray)):
         # fabric array of arrays: return array attribute as a regular array
         return array(dtype=dtype, ndim=ndim)
@@ -6408,8 +6892,18 @@ 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},
-        defaults={"j": None, "k": None},
+        input_types={
+            "arr": array_type(dtype=Any),
+            "i": Any,
+            "j": Any,
+            "k": Any,
+            "l": Any,
+        },
+        defaults={
+            "j": None,
+            "k": None,
+            "l": None,
+        },
         constraint=sametypes,
         hidden=True,
         value_func=view_value_func,
@@ -6513,7 +7007,7 @@ add_builtin(
     hidden=True,
     skip_replay=True,
     group="Utility",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -6530,7 +7024,7 @@ add_builtin(
     dispatch_func=load_dispatch_func,
     hidden=True,
     group="Utility",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -6606,6 +7100,13 @@ def create_atomic_op_value_func(op: str):
                     f"atomic_{op}() operations only work on arrays with [u]int32, [u]int64, float32, or float64 "
                     f"as the underlying scalar types, but got {type_repr(arr_type.dtype)} (with scalar type {type_repr(scalar_type)})"
                 )
+        elif op in ("and", "or", "xor"):
+            supported_atomic_types = (warp.int32, warp.int64, warp.uint32, warp.uint64)
+            if not any(types_equal(scalar_type, x, match_generic=True) for x in supported_atomic_types):
+                raise RuntimeError(
+                    f"atomic_{op}() operations only work on arrays with [u]int32 or [u]int64 "
+                    f"as the underlying scalar types, but got {type_repr(arr_type.dtype)} (with scalar type {type_repr(scalar_type)})"
+                )
         else:
             raise NotImplementedError
 
@@ -6639,7 +7140,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("add"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
-        This function is automatically invoked when using the syntax ``arr[i] += value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i] += value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6651,7 +7153,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("add"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
-        This function is automatically invoked when using the syntax ``arr[i,j] += value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j] += value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6663,7 +7166,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("add"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
-        This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6675,7 +7179,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("add"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
-        This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6688,7 +7193,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("sub"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
-        This function is automatically invoked when using the syntax ``arr[i] -= value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i] -= value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6700,7 +7206,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("sub"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
-        This function is automatically invoked when using the syntax ``arr[i,j] -= value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j] -= value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6712,7 +7219,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("sub"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
-        This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6724,7 +7232,8 @@ for array_type in array_types:
         value_func=create_atomic_op_value_func("sub"),
         dispatch_func=atomic_op_dispatch_func,
         doc="""Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
-        This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.""",
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.""",
         group="Utility",
         skip_replay=True,
     )
@@ -6847,7 +7356,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_cas",
@@ -6861,7 +7370,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_cas",
@@ -6875,7 +7384,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_cas",
@@ -6897,7 +7406,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,
+        is_differentiable=False,
     )
 
     add_builtin(
@@ -6912,7 +7421,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_exch",
@@ -6926,7 +7435,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_exch",
@@ -6940,7 +7449,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,
+        is_differentiable=False,
     )
     add_builtin(
         "atomic_exch",
@@ -6956,6 +7465,177 @@ for array_type in array_types:
         skip_replay=True,
     )
 
+    add_builtin(
+        "atomic_and",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("and"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise AND between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i] &= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_and",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("and"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise AND between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j] &= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_and",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("and"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k] &= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_and",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("and"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise AND between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] &= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+
+    add_builtin(
+        "atomic_or",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("or"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise OR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i] |= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_or",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("or"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise OR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j] |= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_or",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("or"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k] |= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_or",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("or"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise OR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] |= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+
+    add_builtin(
+        "atomic_xor",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("xor"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise XOR between ``value`` and ``arr[i]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i] ^= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_xor",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("xor"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise XOR between ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j] ^= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_xor",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("xor"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k] ^= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+    add_builtin(
+        "atomic_xor",
+        hidden=hidden,
+        input_types={"arr": array_type(dtype=Any), "i": Int, "j": Int, "k": Int, "l": Int, "value": Any},
+        constraint=atomic_op_constraint,
+        value_func=create_atomic_op_value_func("xor"),
+        dispatch_func=atomic_op_dispatch_func,
+        doc="""Atomically performs a bitwise XOR between ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
+
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] ^= value``.""",
+        group="Utility",
+        skip_replay=True,
+        is_differentiable=False,
+    )
+
 
 # used to index into builtin types, i.e.: y = vec3[1]
 def vector_extract_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
@@ -7104,7 +7784,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &quaternion[index]
 add_builtin(
@@ -7115,7 +7795,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &transformation[index]
 add_builtin(
@@ -7126,7 +7806,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &(*vector)[index]
 add_builtin(
@@ -7137,7 +7817,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &(*matrix)[i, j]
 add_builtin(
@@ -7148,7 +7828,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &(*quaternion)[index]
 add_builtin(
@@ -7159,7 +7839,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 # implements &(*transformation)[index]
 add_builtin(
@@ -7170,7 +7850,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -7366,6 +8046,43 @@ add_builtin(
 )
 
 
+# implements vector[idx] &= scalar
+add_builtin(
+    "bit_and_inplace",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    dispatch_func=vector_assign_dispatch_func,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+# implements vector[idx] |= scalar
+add_builtin(
+    "bit_or_inplace",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    dispatch_func=vector_assign_dispatch_func,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+# implements vector[idx] ^= scalar
+add_builtin(
+    "bit_xor_inplace",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    dispatch_func=vector_assign_dispatch_func,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
 def matrix_index_row_value_func(arg_types: Mapping[str, type], arg_values: Mapping[str, Any]):
     mat_type = arg_types["a"]
     row_type = mat_type._wp_row_type_
@@ -7381,7 +8098,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -7400,7 +8117,7 @@ add_builtin(
     hidden=True,
     group="Utility",
     skip_replay=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -7600,6 +8317,78 @@ add_builtin(
 )
 
 
+# implements matrix[i] &= value
+add_builtin(
+    "bit_and_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
+# implements matrix[i,j] &= value
+add_builtin(
+    "bit_and_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "j": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
+# implements matrix[i] |= value
+add_builtin(
+    "bit_or_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
+# implements matrix[i,j] |= value
+add_builtin(
+    "bit_or_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "j": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
+# implements matrix[i] ^= value
+add_builtin(
+    "bit_xor_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
+# implements matrix[i,j] ^= value
+add_builtin(
+    "bit_xor_inplace",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": Any, "j": Any, "value": Any},
+    value_type=None,
+    hidden=True,
+    export=False,
+    group="Utility",
+    is_differentiable=False,
+)
+
+
 for t in scalar_types + vector_types + (bool,):
     if "vec" in t.__name__ or "mat" in t.__name__:
         continue
@@ -7611,7 +8400,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,
+        is_differentiable=False,
     )
 
     add_builtin(
@@ -7622,7 +8411,7 @@ for t in scalar_types + vector_types + (bool,):
         group="Utility",
         hidden=True,
         export=False,
-        missing_grad=True,
+        is_differentiable=False,
     )
 
 
@@ -7641,7 +8430,7 @@ add_builtin(
     doc="Prints an error to stdout if ``a`` and ``b`` are not equal",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "expect_neq",
@@ -7652,7 +8441,7 @@ add_builtin(
     group="Utility",
     hidden=True,
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -7663,7 +8452,7 @@ add_builtin(
     doc="Prints an error to stdout if ``a`` and ``b`` are not equal",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 add_builtin(
     "expect_neq",
@@ -7674,7 +8463,7 @@ add_builtin(
     group="Utility",
     hidden=True,
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -7765,7 +8554,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,
+    is_differentiable=False,
 )
 add_builtin(
     "expect_near",
@@ -7775,7 +8564,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,
+    is_differentiable=False,
 )
 add_builtin(
     "expect_near",
@@ -7785,7 +8574,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,
+    is_differentiable=False,
 )
 add_builtin(
     "expect_near",
@@ -7799,7 +8588,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,
+    is_differentiable=False,
 )
 
 # ---------------------------------
@@ -7810,7 +8599,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -7818,7 +8607,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,
+    is_differentiable=False,
 )
 
 # ---------------------------------
@@ -7899,31 +8688,153 @@ add_builtin(
     input_types={"a": Int, "b": Int},
     value_func=sametypes_create_value_func(Int),
     group="Operators",
-    missing_grad=True,
+    is_differentiable=False,
 )
+add_builtin(
+    "bit_and",
+    input_types={"a": vector(length=Any, dtype=Int), "b": vector(length=Any, dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "bit_and",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int), "b": matrix(shape=(Any, Any), dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+
 add_builtin(
     "bit_or",
     input_types={"a": Int, "b": Int},
     value_func=sametypes_create_value_func(Int),
     group="Operators",
-    missing_grad=True,
+    is_differentiable=False,
 )
+add_builtin(
+    "bit_or",
+    input_types={"a": vector(length=Any, dtype=Int), "b": vector(length=Any, dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "bit_or",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int), "b": matrix(shape=(Any, Any), dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+
 add_builtin(
     "bit_xor",
     input_types={"a": Int, "b": Int},
     value_func=sametypes_create_value_func(Int),
     group="Operators",
-    missing_grad=True,
+    is_differentiable=False,
+)
+add_builtin(
+    "bit_xor",
+    input_types={"a": vector(length=Any, dtype=Int), "b": vector(length=Any, dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "bit_xor",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int), "b": matrix(shape=(Any, Any), dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+
+add_builtin(
+    "lshift",
+    input_types={"a": Int, "b": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "lshift",
+    input_types={"a": vector(length=Any, dtype=Int), "b": vector(length=Any, dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
 )
-add_builtin("lshift", input_types={"a": Int, "b": Int}, value_func=sametypes_create_value_func(Int), group="Operators")
+add_builtin(
+    "lshift",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int), "b": matrix(shape=(Any, Any), dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+
 add_builtin(
     "rshift",
     input_types={"a": Int, "b": Int},
     value_func=sametypes_create_value_func(Int),
     group="Operators",
-    missing_grad=True,
+    is_differentiable=False,
+)
+add_builtin(
+    "rshift",
+    input_types={"a": vector(length=Any, dtype=Int), "b": vector(length=Any, dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "rshift",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int), "b": matrix(shape=(Any, Any), dtype=Int)},
+    constraint=sametypes,
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    doc="",
+    group="Operators",
+    is_differentiable=False,
+)
+
+add_builtin(
+    "invert",
+    input_types={"a": Int},
+    value_func=sametypes_create_value_func(Int),
+    group="Operators",
+    is_differentiable=False,
+)
+add_builtin(
+    "invert",
+    input_types={"a": vector(length=Any, dtype=Int)},
+    value_func=sametypes_create_value_func(vector(length=Any, dtype=Int)),
+    group="Operators",
+    is_differentiable=False,
 )
-add_builtin("invert", input_types={"a": Int}, value_func=sametypes_create_value_func(Int), group="Operators")
+add_builtin(
+    "invert",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Int)},
+    value_func=sametypes_create_value_func(matrix(shape=(Any, Any), dtype=Int)),
+    group="Operators",
+    is_differentiable=False,
+)
+
 
 add_builtin(
     "mul", input_types={"a": Scalar, "b": Scalar}, value_func=sametypes_create_value_func(Scalar), group="Operators"
@@ -8123,7 +9034,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,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -8183,7 +9094,7 @@ add_builtin(
     value_func=sametypes_create_value_func(Scalar),
     doc="",
     group="Operators",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin("pos", input_types={"x": Scalar}, value_func=sametypes_create_value_func(Scalar), group="Operators")
@@ -8232,14 +9143,26 @@ add_builtin(
 )
 
 add_builtin(
-    "unot", input_types={"a": builtins.bool}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True
+    "unot",
+    input_types={"a": builtins.bool},
+    value_type=builtins.bool,
+    doc="",
+    group="Operators",
+    is_differentiable=False,
 )
 for t in int_types:
-    add_builtin("unot", input_types={"a": t}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True)
+    add_builtin(
+        "unot", input_types={"a": t}, value_type=builtins.bool, doc="", group="Operators", is_differentiable=False
+    )
 
 
 add_builtin(
-    "unot", input_types={"a": array(dtype=Any)}, value_type=builtins.bool, doc="", group="Operators", missing_grad=True
+    "unot",
+    input_types={"a": array(dtype=Any)},
+    value_type=builtins.bool,
+    doc="",
+    group="Operators",
+    is_differentiable=False,
 )
 
 
@@ -8312,6 +9235,45 @@ add_builtin(
     export=False,
 )
 
+add_builtin(
+    "bit_and",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_func=tile_binary_map_value_func,
+    # dispatch_func=tile_map_dispatch_func,
+    # variadic=True,
+    native_func="tile_bit_and",
+    doc="Bitwise AND each element of two tiles together",
+    group="Tile Primitives",
+    export=False,
+    is_differentiable=False,
+)
+
+add_builtin(
+    "bit_or",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_func=tile_binary_map_value_func,
+    # dispatch_func=tile_map_dispatch_func,
+    # variadic=True,
+    native_func="tile_bit_or",
+    doc="Bitwise OR each element of two tiles together",
+    group="Tile Primitives",
+    export=False,
+    is_differentiable=False,
+)
+
+add_builtin(
+    "bit_xor",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_func=tile_binary_map_value_func,
+    # dispatch_func=tile_map_dispatch_func,
+    # variadic=True,
+    native_func="tile_bit_xor",
+    doc="Bitwise XOR each element of two tiles together",
+    group="Tile Primitives",
+    export=False,
+    is_differentiable=False,
+)
+
 
 add_builtin(
     "mul",
@@ -8373,6 +9335,45 @@ add_builtin(
 )
 
 
+add_builtin(
+    "bit_and_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_type=None,
+    dispatch_func=tile_inplace_dispatch_func,
+    export=False,
+    hidden=True,
+    native_func="tile_bit_and_inplace",
+    group="Operators",
+    is_differentiable=False,
+)
+
+
+add_builtin(
+    "bit_or_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_type=None,
+    dispatch_func=tile_inplace_dispatch_func,
+    export=False,
+    hidden=True,
+    native_func="tile_bit_or_inplace",
+    group="Operators",
+    is_differentiable=False,
+)
+
+
+add_builtin(
+    "bit_xor_inplace",
+    input_types={"a": tile(dtype=Any, shape=Tuple[int, ...]), "b": tile(dtype=Any, shape=Tuple[int, ...])},
+    value_type=None,
+    dispatch_func=tile_inplace_dispatch_func,
+    export=False,
+    hidden=True,
+    native_func="tile_bit_xor_inplace",
+    group="Operators",
+    is_differentiable=False,
+)
+
+
 def tile_diag_add_value_func(arg_types, arg_values):
     if arg_types is None:
         return tile(dtype=Any, shape=Tuple[int, int])
@@ -8414,7 +9415,7 @@ def tile_diag_add_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     a = arg_values["a"]
     d = arg_values["d"]
@@ -8434,7 +9435,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,
+    is_differentiable=False,
 )
 
 
@@ -8491,7 +9492,7 @@ def tile_matmul_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     a = arg_values["a"]
     b = arg_values["b"]
@@ -8529,7 +9530,7 @@ def tile_matmul_lto_dispatch_func(
     num_threads = options["block_dim"]
     arch = options["output_arch"]
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ((0, 0, 0, a, b, out), template_args, [], 0)
     else:
@@ -8542,7 +9543,7 @@ def tile_matmul_lto_dispatch_func(
 
         # generate the LTOs
         #    C += A * B
-        (fun_forward, lto_forward) = warp.build.build_lto_dot(
+        (fun_forward, lto_forward) = warp._src.build.build_lto_dot(
             M,
             N,
             K,
@@ -8558,7 +9559,7 @@ def tile_matmul_lto_dispatch_func(
         )
         if warp.config.enable_backward:
             # adjA += adjC * B^T - Transpose ~= flipped layout
-            (fun_backward_A, lto_backward_A) = warp.build.build_lto_dot(
+            (fun_backward_A, lto_backward_A) = warp._src.build.build_lto_dot(
                 M,
                 K,
                 N,
@@ -8573,7 +9574,7 @@ def tile_matmul_lto_dispatch_func(
                 builder,
             )
             # adjB += A^T * adjC - Transpose ~= flipped layout
-            (fun_backward_B, lto_backward_B) = warp.build.build_lto_dot(
+            (fun_backward_B, lto_backward_B) = warp._src.build.build_lto_dot(
                 K,
                 N,
                 M,
@@ -8690,7 +9691,7 @@ def tile_fft_generic_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
     direction: str | None = None,
 ):
     inout = arg_values["inout"]
@@ -8719,12 +9720,12 @@ def tile_fft_generic_lto_dispatch_func(
     arch = options["output_arch"]
     ept = size // num_threads
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ([], [], [], 0)
     else:
         # generate the LTO
-        lto_symbol, lto_code_data, shared_memory_bytes = warp.build.build_lto_fft(
+        lto_symbol, lto_code_data, shared_memory_bytes = warp._src.build.build_lto_fft(
             arch, size, ept, direction, dir, precision, builder
         )
 
@@ -8762,7 +9763,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -8784,7 +9785,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -8829,7 +9830,7 @@ def tile_cholesky_generic_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     a = arg_values["A"]
     # force source tile to shared memory
@@ -8849,7 +9850,7 @@ def tile_cholesky_generic_lto_dispatch_func(
 
     arch = options["output_arch"]
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ((0, a, out), [], [], 0)
     else:
@@ -8864,7 +9865,7 @@ def tile_cholesky_generic_lto_dispatch_func(
         req_smem_bytes = a.type.size * type_size_in_bytes(a.type.dtype)
 
         # generate the LTO
-        lto_symbol, lto_code_data = warp.build.build_lto_solver(
+        lto_symbol, lto_code_data = warp._src.build.build_lto_solver(
             M,
             N,
             1,
@@ -8909,7 +9910,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -8953,7 +9954,7 @@ def tile_cholesky_solve_generic_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     L = arg_values["L"]
     y = arg_values["y"]
@@ -8982,7 +9983,7 @@ def tile_cholesky_solve_generic_lto_dispatch_func(
 
     arch = options["output_arch"]
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ((0, L, y, x), [], [], 0)
     else:
@@ -8998,7 +9999,7 @@ def tile_cholesky_solve_generic_lto_dispatch_func(
         req_smem_bytes = (x.type.size + y.type.size + L.type.size) * type_size_in_bytes(L.type.dtype)
 
         # generate the LTO
-        lto_symbol, lto_code_data = warp.build.build_lto_solver(
+        lto_symbol, lto_code_data = warp._src.build.build_lto_solver(
             M,
             N,
             NRHS,
@@ -9040,7 +10041,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -9050,7 +10051,7 @@ def tile_lower_solve_generic_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     L = arg_values["L"]
     y = arg_values["y"]
@@ -9079,7 +10080,7 @@ def tile_lower_solve_generic_lto_dispatch_func(
 
     arch = options["output_arch"]
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ((0, L, y, z), [], [], 0)
     else:
@@ -9095,7 +10096,7 @@ def tile_lower_solve_generic_lto_dispatch_func(
         req_smem_bytes = (z.type.size + y.type.size + L.type.size) * type_size_in_bytes(L.type.dtype)
 
         # generate the LTO
-        lto_symbol, lto_code_data = warp.build.build_lto_solver(
+        lto_symbol, lto_code_data = warp._src.build.build_lto_solver(
             M,
             N,
             NRHS,
@@ -9173,7 +10174,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -9183,7 +10184,7 @@ def tile_upper_solve_generic_lto_dispatch_func(
     return_values: List[Var],
     arg_values: Mapping[str, Var],
     options: Mapping[str, Any],
-    builder: warp.context.ModuleBuilder,
+    builder: warp._src.context.ModuleBuilder,
 ):
     U = arg_values["U"]
     z = arg_values["z"]
@@ -9212,7 +10213,7 @@ def tile_upper_solve_generic_lto_dispatch_func(
 
     arch = options["output_arch"]
 
-    if arch is None or not warp.context.runtime.core.wp_is_mathdx_enabled():
+    if arch is None or not warp._src.context.runtime.core.wp_is_mathdx_enabled():
         # CPU/no-MathDx dispatch
         return ((0, U, z, x), [], [], 0)
     else:
@@ -9228,7 +10229,7 @@ def tile_upper_solve_generic_lto_dispatch_func(
         req_smem_bytes = (x.type.size + z.type.size + U.type.size) * type_size_in_bytes(U.type.dtype)
 
         # generate the LTO
-        lto_symbol, lto_code_data = warp.build.build_lto_solver(
+        lto_symbol, lto_code_data = warp._src.build.build_lto_solver(
             M,
             N,
             NRHS,
@@ -9306,7 +10307,7 @@ add_builtin(
     group="Tile Primitives",
     export=False,
     namespace="",
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -9326,7 +10327,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,
+    is_differentiable=False,
 )
 
 
@@ -9351,7 +10352,7 @@ add_builtin(
     doc="Return the number of elements in a vector.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -9361,7 +10362,7 @@ add_builtin(
     doc="Return the number of elements in a quaternion.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -9371,7 +10372,7 @@ add_builtin(
     doc="Return the number of rows in a matrix.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -9381,7 +10382,7 @@ add_builtin(
     doc="Return the number of elements in a transformation.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -9391,7 +10392,7 @@ add_builtin(
     doc="Return the size of the first dimension in an array.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 add_builtin(
@@ -9401,7 +10402,62 @@ add_builtin(
     doc="Return the number of rows in a tile.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
+)
+
+
+def cast_value_func(arg_types, arg_values):
+    # Return generic type for doc builds.
+    if arg_types is None:
+        return Any
+
+    return arg_values["dtype"]
+
+
+def cast_dispatch_func(input_types: Mapping[str, type], return_type: Any, args: Mapping[str, Var]):
+    func_args = (args["a"],)
+    template_args = (args["dtype"],)
+    return (func_args, template_args)
+
+
+add_builtin(
+    "cast",
+    input_types={"a": Any, "dtype": Any},
+    value_func=cast_value_func,
+    dispatch_func=cast_dispatch_func,
+    group="Utility",
+    export=False,
+    is_differentiable=False,
+    doc="""Reinterpret a value as a different type while preserving its bit pattern.
+
+    :param a: The value to cast
+    :param dtype: The target type
+
+    Example:
+
+    .. code-block:: python
+
+        @wp.struct
+        class MyStruct:
+            f: wp.float16
+            i: wp.int16
+
+
+        @wp.kernel
+        def compute():
+            x = wp.int32(0x40000000)
+            x_casted = wp.cast(x, wp.float32)
+            wp.expect_eq(x_casted, 2.0) # 0x40000000
+
+            s = MyStruct()
+            s.f = wp.float16(2.0) # 0x4000
+            s.i = wp.int16(4096) # 0x1000
+            s_casted = wp.cast(s, wp.int32)
+            wp.expect_eq(s_casted, 0x10004000)
+
+
+        wp.launch(compute, dim=1)
+    """,
 )
 
 
@@ -9428,7 +10484,7 @@ add_builtin(
     doc="Construct a tuple from a list of values",
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
     export=False,
 )
 
@@ -9465,7 +10521,7 @@ add_builtin(
     dispatch_func=tuple_extract_dispatch_func,
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 
@@ -9476,7 +10532,7 @@ add_builtin(
     doc="Return the number of elements in a tuple.",
     group="Utility",
     export=False,
-    missing_grad=True,
+    is_differentiable=False,
 )
 
 # ---------------------------------
@@ -9495,5 +10551,5 @@ add_builtin(
     export=False,
     group="Utility",
     hidden=True,
-    missing_grad=True,
+    is_differentiable=False,
 )