warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/builtins.py

@@ -5,24 +5,22 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-from .context import add_builtin
+import builtins
+from typing import Any, Callable, Tuple
 
+from warp.codegen import Reference
 from warp.types import *
 
-from typing import Tuple
-from typing import List
-from typing import Dict
-from typing import Any
-from typing import Callable
+from .context import add_builtin
 
 
 def sametype_value_func(default):
-    def fn(args, kwds, _):
-        if args is None:
+    def fn(arg_types, kwds, _):
+        if arg_types is None:
             return default
-        if not all(types_equal(args[0].type, a.type) for a in args[1:]):
-            raise RuntimeError(f"Input types must be the same, found: {[type_repr(a.type) for a in args]}")
-        return args[0].type
+        if not all(types_equal(arg_types[0], t) for t in arg_types[1:]):
+            raise RuntimeError(f"Input types must be the same, found: {[type_repr(t) for t in arg_types]}")
+        return arg_types[0]
 
     return fn
 
@@ -50,7 +48,7 @@ add_builtin(
     "clamp",
     input_types={"x": Scalar, "a": Scalar, "b": Scalar},
     value_func=sametype_value_func(Scalar),
-    doc="Clamp the value of x to the range [a, b].",
+    doc="Clamp the value of ``x`` to the range [a, b].",
     group="Scalar Math",
 )
 
@@ -58,14 +56,14 @@ add_builtin(
     "abs",
     input_types={"x": Scalar},
     value_func=sametype_value_func(Scalar),
-    doc="Return the absolute value of x.",
+    doc="Return the absolute value of ``x``.",
     group="Scalar Math",
 )
 add_builtin(
     "sign",
     input_types={"x": Scalar},
     value_func=sametype_value_func(Scalar),
-    doc="Return -1 if x < 0, return 1 otherwise.",
+    doc="Return -1 if ``x`` < 0, return 1 otherwise.",
     group="Scalar Math",
 )
 
@@ -73,14 +71,14 @@ add_builtin(
     "step",
     input_types={"x": Scalar},
     value_func=sametype_value_func(Scalar),
-    doc="Return 1.0 if x < 0.0, return 0.0 otherwise.",
+    doc="Return 1.0 if ``x`` < 0.0, return 0.0 otherwise.",
     group="Scalar Math",
 )
 add_builtin(
     "nonzero",
     input_types={"x": Scalar},
     value_func=sametype_value_func(Scalar),
-    doc="Return 1.0 if x is not equal to zero, return 0.0 otherwise.",
+    doc="Return 1.0 if ``x`` is not equal to zero, return 0.0 otherwise.",
     group="Scalar Math",
 )
 
@@ -88,91 +86,101 @@ add_builtin(
     "sin",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the sine of x in radians.",
+    doc="Return the sine of ``x`` in radians.",
     group="Scalar Math",
 )
 add_builtin(
     "cos",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the cosine of x in radians.",
+    doc="Return the cosine of ``x`` in radians.",
     group="Scalar Math",
 )
 add_builtin(
     "acos",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return arccos of x in radians. Inputs are automatically clamped to [-1.0, 1.0].",
+    doc="Return arccos of ``x`` in radians. Inputs are automatically clamped to [-1.0, 1.0].",
     group="Scalar Math",
 )
 add_builtin(
     "asin",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return arcsin of x in radians. Inputs are automatically clamped to [-1.0, 1.0].",
+    doc="Return arcsin of ``x`` in radians. Inputs are automatically clamped to [-1.0, 1.0].",
     group="Scalar Math",
 )
 add_builtin(
     "sqrt",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the sqrt of x, where x is positive.",
+    doc="Return the square root of ``x``, where ``x`` is positive.",
     group="Scalar Math",
+    require_original_output_arg=True,
+)
+add_builtin(
+    "cbrt",
+    input_types={"x": Float},
+    value_func=sametype_value_func(Float),
+    doc="Return the cube root of ``x``.",
+    group="Scalar Math",
+    require_original_output_arg=True,
 )
 add_builtin(
     "tan",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return tangent of x in radians.",
+    doc="Return the tangent of ``x`` in radians.",
     group="Scalar Math",
 )
 add_builtin(
     "atan",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return arctan of x.",
+    doc="Return the arctangent of ``x`` in radians.",
     group="Scalar Math",
 )
 add_builtin(
     "atan2",
     input_types={"y": Float, "x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return atan2 of x.",
+    doc="Return the 2-argument arctangent, atan2, of the point ``(x, y)`` in radians.",
     group="Scalar Math",
 )
 add_builtin(
     "sinh",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the sinh of x.",
+    doc="Return the sinh of ``x``.",
     group="Scalar Math",
 )
 add_builtin(
     "cosh",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the cosh of x.",
+    doc="Return the cosh of ``x``.",
     group="Scalar Math",
 )
 add_builtin(
     "tanh",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the tanh of x.",
+    doc="Return the tanh of ``x``.",
     group="Scalar Math",
+    require_original_output_arg=True,
 )
 add_builtin(
     "degrees",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Convert radians into degrees.",
+    doc="Convert ``x`` from radians into degrees.",
     group="Scalar Math",
 )
 add_builtin(
     "radians",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Convert degrees into radians.",
+    doc="Convert ``x`` from degrees into radians.",
     group="Scalar Math",
 )
 
@@ -180,36 +188,38 @@ add_builtin(
     "log",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the natural log (base-e) of x, where x is positive.",
+    doc="Return the natural logarithm (base-e) of ``x``, where ``x`` is positive.",
     group="Scalar Math",
 )
 add_builtin(
     "log2",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the natural log (base-2) of x, where x is positive.",
+    doc="Return the binary logarithm (base-2) of ``x``, where ``x`` is positive.",
     group="Scalar Math",
 )
 add_builtin(
     "log10",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the natural log (base-10) of x, where x is positive.",
+    doc="Return the common logarithm (base-10) of ``x``, where ``x`` is positive.",
     group="Scalar Math",
 )
 add_builtin(
     "exp",
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
-    doc="Return base-e exponential, e^x.",
+    doc="Return the value of the exponential function :math:`e^x`.",
     group="Scalar Math",
+    require_original_output_arg=True,
 )
 add_builtin(
     "pow",
     input_types={"x": Float, "y": Float},
     value_func=sametype_value_func(Float),
-    doc="Return the result of x raised to power of y.",
+    doc="Return the result of ``x`` raised to power of ``y``.",
     group="Scalar Math",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -217,9 +227,9 @@ add_builtin(
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
     group="Scalar Math",
-    doc="""Calculate the nearest integer value, rounding halfway cases away from zero.
-    This is the most intuitive form of rounding in the colloquial sense, but can be slower than other options like ``warp.rint()``.
-    Differs from ``numpy.round()``, which behaves the same way as ``numpy.rint()``.""",
+    doc="""Return the nearest integer value to ``x``, rounding halfway cases away from zero.
+   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()`.""",
 )
 
 add_builtin(
@@ -227,9 +237,8 @@ add_builtin(
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
     group="Scalar Math",
-    doc="""Calculate the nearest integer value, rounding halfway cases to nearest even integer.
-    It is generally faster than ``warp.round()``.
-    Equivalent to ``numpy.rint()``.""",
+    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()`.""",
 )
 
 add_builtin(
@@ -237,10 +246,10 @@ add_builtin(
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
     group="Scalar Math",
-    doc="""Calculate the nearest integer that is closer to zero than x.
-    In other words, it discards the fractional part of x.
-    It is similar to casting ``float(int(x))``, but preserves the negative sign when x is in the range [-0.0, -1.0).
-    Equivalent to ``numpy.trunc()`` and ``numpy.fix()``.""",
+    doc="""Return the nearest integer that is closer to zero than ``x``.
+   In other words, it discards the fractional part of ``x``.
+   It is similar to casting ``float(int(x))``, but preserves the negative sign when x is in the range [-0.0, -1.0).
+   Equivalent to :func:`numpy.trunc()` and :func:`numpy.fix()`.""",
 )
 
 add_builtin(
@@ -248,7 +257,7 @@ add_builtin(
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
     group="Scalar Math",
-    doc="""Calculate the largest integer that is less than or equal to x.""",
+    doc="""Return the largest integer that is less than or equal to ``x``.""",
 )
 
 add_builtin(
@@ -256,22 +265,31 @@ add_builtin(
     input_types={"x": Float},
     value_func=sametype_value_func(Float),
     group="Scalar Math",
-    doc="""Calculate the smallest integer that is greater than or equal to x.""",
+    doc="""Return the smallest integer that is greater than or equal to ``x``.""",
+)
+
+add_builtin(
+    "frac",
+    input_types={"x": Float},
+    value_func=sametype_value_func(Float),
+    group="Scalar Math",
+    doc="""Retrieve the fractional part of x.
+    In other words, it discards the integer part of x and is equivalent to ``x - trunc(x)``.""",
 )
 
 
-def infer_scalar_type(args):
-    if args is None:
+def infer_scalar_type(arg_types):
+    if arg_types is None:
         return Scalar
 
-    def iterate_scalar_types(args):
-        for a in args:
-            if hasattr(a.type, "_wp_scalar_type_"):
-                yield a.type._wp_scalar_type_
-            elif a.type in scalar_types:
-                yield a.type
+    def iterate_scalar_types(arg_types):
+        for t in arg_types:
+            if hasattr(t, "_wp_scalar_type_"):
+                yield t._wp_scalar_type_
+            elif t in scalar_types:
+                yield t
 
-    scalarTypes = set(iterate_scalar_types(args))
+    scalarTypes = set(iterate_scalar_types(arg_types))
     if len(scalarTypes) > 1:
         raise RuntimeError(
             f"Couldn't figure out return type as arguments have multiple precisions: {list(scalarTypes)}"
@@ -279,13 +297,13 @@ def infer_scalar_type(args):
     return list(scalarTypes)[0]
 
 
-def sametype_scalar_value_func(args, kwds, _):
-    if args is None:
+def sametype_scalar_value_func(arg_types, kwds, _):
+    if arg_types is None:
         return Scalar
-    if not all(types_equal(args[0].type, a.type) for a in args[1:]):
-        raise RuntimeError(f"Input types must be exactly the same, {[a.type for a in args]}")
+    if not all(types_equal(arg_types[0], t) for t in arg_types[1:]):
+        raise RuntimeError(f"Input types must be exactly the same, {[t for t in arg_types]}")
 
-    return infer_scalar_type(args)
+    return infer_scalar_type(arg_types)
 
 
 # ---------------------------------
@@ -310,14 +328,14 @@ add_builtin(
     "min",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
     value_func=sametype_value_func(vector(length=Any, dtype=Scalar)),
-    doc="Return the element wise minimum of two vectors.",
+    doc="Return the element-wise minimum of two vectors.",
     group="Vector Math",
 )
 add_builtin(
     "max",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
     value_func=sametype_value_func(vector(length=Any, dtype=Scalar)),
-    doc="Return the element wise maximum of two vectors.",
+    doc="Return the element-wise maximum of two vectors.",
     group="Vector Math",
 )
 
@@ -325,41 +343,41 @@ add_builtin(
     "min",
     input_types={"v": vector(length=Any, dtype=Scalar)},
     value_func=sametype_scalar_value_func,
-    doc="Return the minimum element of a vector.",
+    doc="Return the minimum element of a vector ``v``.",
     group="Vector Math",
 )
 add_builtin(
     "max",
     input_types={"v": vector(length=Any, dtype=Scalar)},
     value_func=sametype_scalar_value_func,
-    doc="Return the maximum element of a vector.",
+    doc="Return the maximum element of a vector ``v``.",
     group="Vector Math",
 )
 
 add_builtin(
     "argmin",
     input_types={"v": vector(length=Any, dtype=Scalar)},
-    value_func=lambda args, kwds, _: warp.uint32,
-    doc="Return the index of the minimum element of a vector.",
+    value_func=lambda arg_types, kwds, _: warp.uint32,
+    doc="Return the index of the minimum element of a vector ``v``.",
     group="Vector Math",
     missing_grad=True,
 )
 add_builtin(
     "argmax",
     input_types={"v": vector(length=Any, dtype=Scalar)},
-    value_func=lambda args, kwds, _: warp.uint32,
-    doc="Return the index of the maximum element of a vector.",
+    value_func=lambda arg_types, kwds, _: warp.uint32,
+    doc="Return the index of the maximum element of a vector ``v``.",
     group="Vector Math",
     missing_grad=True,
 )
 
 
-def value_func_outer(args, kwds, _):
-    if args is None:
+def value_func_outer(arg_types, kwds, _):
+    if arg_types is None:
         return matrix(shape=(Any, Any), dtype=Scalar)
 
-    scalarType = infer_scalar_type(args)
-    vectorLengths = [i.type._length_ for i in args]
+    scalarType = infer_scalar_type(arg_types)
+    vectorLengths = [t._length_ for t in arg_types]
     return matrix(shape=(vectorLengths), dtype=scalarType)
 
 
@@ -368,7 +386,7 @@ add_builtin(
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
     value_func=value_func_outer,
     group="Vector Math",
-    doc="Compute the outer product x*y^T for two vec2 objects.",
+    doc="Compute the outer product ``x*y^T`` for two vectors.",
 )
 
 add_builtin(
@@ -376,14 +394,14 @@ add_builtin(
     input_types={"x": vector(length=3, dtype=Scalar), "y": vector(length=3, dtype=Scalar)},
     value_func=sametype_value_func(vector(length=3, dtype=Scalar)),
     group="Vector Math",
-    doc="Compute the cross product of two 3d vectors.",
+    doc="Compute the cross product of two 3D vectors.",
 )
 add_builtin(
     "skew",
     input_types={"x": vector(length=3, dtype=Scalar)},
-    value_func=lambda args, kwds, _: matrix(shape=(3, 3), dtype=args[0].type._wp_scalar_type_),
+    value_func=lambda arg_types, kwds, _: matrix(shape=(3, 3), dtype=arg_types[0]._wp_scalar_type_),
     group="Vector Math",
-    doc="Compute the skew symmetric matrix for a 3d vector.",
+    doc="Compute the skew-symmetric 3x3 matrix for a 3D vector ``x``.",
 )
 
 add_builtin(
@@ -391,59 +409,62 @@ add_builtin(
     input_types={"x": vector(length=Any, dtype=Float)},
     value_func=sametype_scalar_value_func,
     group="Vector Math",
-    doc="Compute the length of a vector.",
+    doc="Compute the length of a vector ``x``.",
+    require_original_output_arg=True,
 )
 add_builtin(
     "length",
     input_types={"x": quaternion(dtype=Float)},
     value_func=sametype_scalar_value_func,
     group="Vector Math",
-    doc="Compute the length of a quaternion.",
+    doc="Compute the length of a quaternion ``x``.",
+    require_original_output_arg=True,
 )
 add_builtin(
     "length_sq",
     input_types={"x": vector(length=Any, dtype=Scalar)},
     value_func=sametype_scalar_value_func,
     group="Vector Math",
-    doc="Compute the squared length of a 2d vector.",
+    doc="Compute the squared length of a 2D vector ``x``.",
 )
 add_builtin(
     "length_sq",
     input_types={"x": quaternion(dtype=Scalar)},
     value_func=sametype_scalar_value_func,
     group="Vector Math",
-    doc="Compute the squared length of a quaternion.",
+    doc="Compute the squared length of a quaternion ``x``.",
 )
 add_builtin(
     "normalize",
     input_types={"x": vector(length=Any, dtype=Float)},
     value_func=sametype_value_func(vector(length=Any, dtype=Scalar)),
     group="Vector Math",
-    doc="Compute the normalized value of x, if length(x) is 0 then the zero vector is returned.",
+    doc="Compute the normalized value of ``x``. If ``length(x)`` is 0 then the zero vector is returned.",
+    require_original_output_arg=True,
 )
 add_builtin(
     "normalize",
     input_types={"x": quaternion(dtype=Float)},
     value_func=sametype_value_func(quaternion(dtype=Scalar)),
     group="Vector Math",
-    doc="Compute the normalized value of x, if length(x) is 0 then the zero quat is returned.",
+    doc="Compute the normalized value of ``x``. If ``length(x)`` is 0, then the zero quaternion is returned.",
 )
 
 add_builtin(
     "transpose",
     input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
-    value_func=lambda args, kwds, _: matrix(
-        shape=(args[0].type._shape_[1], args[0].type._shape_[0]), dtype=args[0].type._wp_scalar_type_
+    value_func=lambda arg_types, kwds, _: matrix(
+        shape=(arg_types[0]._shape_[1], arg_types[0]._shape_[0]), dtype=arg_types[0]._wp_scalar_type_
     ),
     group="Vector Math",
-    doc="Return the transpose of the matrix m",
+    doc="Return the transpose of the matrix ``m``.",
 )
 
 
-def value_func_mat_inv(args, kwds, _):
-    if args is None:
+def value_func_mat_inv(arg_types, kwds, _):
+    if arg_types is None:
         return matrix(shape=(Any, Any), dtype=Float)
-    return args[0].type
+    return arg_types[0]
 
 
 add_builtin(
@@ -451,7 +472,8 @@ add_builtin(
     input_types={"m": matrix(shape=(2, 2), dtype=Float)},
     value_func=value_func_mat_inv,
     group="Vector Math",
-    doc="Return the inverse of a 2x2 matrix m",
+    doc="Return the inverse of a 2x2 matrix ``m``.",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -459,7 +481,8 @@ add_builtin(
     input_types={"m": matrix(shape=(3, 3), dtype=Float)},
     value_func=value_func_mat_inv,
     group="Vector Math",
-    doc="Return the inverse of a 3x3 matrix m",
+    doc="Return the inverse of a 3x3 matrix ``m``.",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -467,14 +490,15 @@ add_builtin(
     input_types={"m": matrix(shape=(4, 4), dtype=Float)},
     value_func=value_func_mat_inv,
     group="Vector Math",
-    doc="Return the inverse of a 4x4 matrix m",
+    doc="Return the inverse of a 4x4 matrix ``m``.",
+    require_original_output_arg=True,
 )
 
 
-def value_func_mat_det(args, kwds, _):
-    if args is None:
+def value_func_mat_det(arg_types, kwds, _):
+    if arg_types is None:
         return Scalar
-    return args[0].type._wp_scalar_type_
+    return arg_types[0]._wp_scalar_type_
 
 
 add_builtin(
@@ -482,7 +506,7 @@ add_builtin(
     input_types={"m": matrix(shape=(2, 2), dtype=Float)},
     value_func=value_func_mat_det,
     group="Vector Math",
-    doc="Return the determinant of a 2x2 matrix m",
+    doc="Return the determinant of a 2x2 matrix ``m``.",
 )
 
 add_builtin(
@@ -490,7 +514,7 @@ add_builtin(
     input_types={"m": matrix(shape=(3, 3), dtype=Float)},
     value_func=value_func_mat_det,
     group="Vector Math",
-    doc="Return the determinant of a 3x3 matrix m",
+    doc="Return the determinant of a 3x3 matrix ``m``.",
 )
 
 add_builtin(
@@ -498,16 +522,16 @@ add_builtin(
     input_types={"m": matrix(shape=(4, 4), dtype=Float)},
     value_func=value_func_mat_det,
     group="Vector Math",
-    doc="Return the determinant of a 4x4 matrix m",
+    doc="Return the determinant of a 4x4 matrix ``m``.",
 )
 
 
-def value_func_mat_trace(args, kwds, _):
-    if args is None:
+def value_func_mat_trace(arg_types, kwds, _):
+    if arg_types is None:
         return Scalar
-    if args[0].type._shape_[0] != args[0].type._shape_[1]:
-        raise RuntimeError(f"Matrix shape is {args[0].type._shape_}. Cannot find the trace of non square matrices")
-    return args[0].type._wp_scalar_type_
+    if arg_types[0]._shape_[0] != arg_types[0]._shape_[1]:
+        raise RuntimeError(f"Matrix shape is {arg_types[0]._shape_}. Cannot find the trace of non square matrices")
+    return arg_types[0]._wp_scalar_type_
 
 
 add_builtin(
@@ -515,15 +539,15 @@ add_builtin(
     input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
     value_func=value_func_mat_trace,
     group="Vector Math",
-    doc="Return the trace of the matrix m",
+    doc="Return the trace of the matrix ``m``.",
 )
 
 
-def value_func_diag(args, kwds, _):
-    if args is None:
+def value_func_diag(arg_types, kwds, _):
+    if arg_types is None:
         return matrix(shape=(Any, Any), dtype=Scalar)
     else:
-        return matrix(shape=(args[0].type._length_, args[0].type._length_), dtype=args[0].type._wp_scalar_type_)
+        return matrix(shape=(arg_types[0]._length_, arg_types[0]._length_), dtype=arg_types[0]._wp_scalar_type_)
 
 
 add_builtin(
@@ -531,7 +555,27 @@ add_builtin(
     input_types={"d": vector(length=Any, dtype=Scalar)},
     value_func=value_func_diag,
     group="Vector Math",
-    doc="Returns a matrix with the components of the vector d on the diagonal",
+    doc="Returns a matrix with the components of the vector ``d`` on the diagonal.",
+)
+
+
+def value_func_get_diag(arg_types, kwds, _):
+    if arg_types is None:
+        return vector(length=(Any), dtype=Scalar)
+    else:
+        if arg_types[0]._shape_[0] != arg_types[0]._shape_[1]:
+            raise RuntimeError(
+                f"Matrix shape is {arg_types[0]._shape_}; get_diag is only available for square matrices."
+            )
+        return vector(length=arg_types[0]._shape_[0], dtype=arg_types[0]._wp_scalar_type_)
+
+
+add_builtin(
+    "get_diag",
+    input_types={"m": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_func=value_func_get_diag,
+    group="Vector Math",
+    doc="Returns a vector containing the diagonal elements of the square matrix ``m``.",
 )
 
 add_builtin(
@@ -539,14 +583,15 @@ add_builtin(
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
     value_func=sametype_value_func(vector(length=Any, dtype=Scalar)),
     group="Vector Math",
-    doc="Component wise multiply of two 2d vectors.",
+    doc="Component-wise multiplication of two 2D vectors.",
 )
 add_builtin(
     "cw_div",
     input_types={"x": vector(length=Any, dtype=Scalar), "y": vector(length=Any, dtype=Scalar)},
     value_func=sametype_value_func(vector(length=Any, dtype=Scalar)),
     group="Vector Math",
-    doc="Component wise division of two 2d vectors.",
+    doc="Component-wise division of two 2D vectors.",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -554,14 +599,15 @@ add_builtin(
     input_types={"x": matrix(shape=(Any, Any), dtype=Scalar), "y": matrix(shape=(Any, Any), dtype=Scalar)},
     value_func=sametype_value_func(matrix(shape=(Any, Any), dtype=Scalar)),
     group="Vector Math",
-    doc="Component wise multiply of two 2d vectors.",
+    doc="Component-wise multiplication of two 2D vectors.",
 )
 add_builtin(
     "cw_div",
     input_types={"x": matrix(shape=(Any, Any), dtype=Scalar), "y": matrix(shape=(Any, Any), dtype=Scalar)},
     value_func=sametype_value_func(matrix(shape=(Any, Any), dtype=Scalar)),
     group="Vector Math",
-    doc="Component wise division of two 2d vectors.",
+    doc="Component-wise division of two 2D vectors.",
+    require_original_output_arg=True,
 )
 
 
@@ -573,16 +619,19 @@ for t in scalar_types_all:
             t.__name__, input_types={"u": u}, value_type=t, doc="", hidden=True, group="Scalar Math", export=False
         )
 
+for u in [bool, builtins.bool]:
+    add_builtin(bool.__name__, input_types={"u": u}, value_type=bool, doc="", hidden=True, export=False, namespace="")
+
 
-def vector_constructor_func(args, kwds, templates):
-    if args is None:
+def vector_constructor_func(arg_types, kwds, templates):
+    if arg_types is None:
         return vector(length=Any, dtype=Scalar)
 
     if templates is None or len(templates) == 0:
         # handle construction of anonymous (undeclared) vector types
 
         if "length" in kwds:
-            if len(args) == 0:
+            if len(arg_types) == 0:
                 if "dtype" not in kwds:
                     raise RuntimeError(
                         "vec() must have dtype as a keyword argument if it has no positional arguments, e.g.: wp.vector(length=5, dtype=wp.float32)"
@@ -592,34 +641,54 @@ def vector_constructor_func(args, kwds, templates):
                 veclen = kwds["length"]
                 vectype = kwds["dtype"]
 
-            elif len(args) == 1:
+            elif len(arg_types) == 1:
                 # value initialization e.g.: wp.vec(1.0, length=5)
                 veclen = kwds["length"]
-                vectype = args[0].type
+                vectype = arg_types[0]
+                if getattr(vectype, "_wp_generic_type_str_", None) == "vec_t":
+                    # constructor from another vector
+                    if vectype._length_ != veclen:
+                        raise RuntimeError(
+                            f"Incompatible vector lengths for casting copy constructor, {veclen} vs {vectype._length_}"
+                        )
+                    vectype = vectype._wp_scalar_type_
             else:
                 raise RuntimeError(
                     "vec() must have one scalar argument or the dtype keyword argument if the length keyword argument is specified, e.g.: wp.vec(1.0, length=5)"
                 )
 
         else:
-            if len(args) == 0:
+            if len(arg_types) == 0:
                 raise RuntimeError(
                     "vec() must have at least one numeric argument, if it's length, dtype is not specified"
                 )
 
             if "dtype" in kwds:
+                # casting constructor
+                if len(arg_types) == 1 and types_equal(
+                    arg_types[0], vector(length=Any, dtype=Scalar), match_generic=True
+                ):
+                    veclen = arg_types[0]._length_
+                    vectype = kwds["dtype"]
+                    templates.append(veclen)
+                    templates.append(vectype)
+                    return vector(length=veclen, dtype=vectype)
                 raise RuntimeError(
                     "vec() should not have dtype specified if numeric arguments are given, the dtype will be inferred from the argument types"
                 )
 
             # component wise construction of an anonymous vector, e.g. wp.vec(wp.float16(1.0), wp.float16(2.0), ....)
             # we infer the length and data type from the number and type of the arg values
-            veclen = len(args)
-            vectype = args[0].type
-
-            if not all(vectype == a.type for a in args):
+            veclen = len(arg_types)
+            vectype = arg_types[0]
+
+            if len(arg_types) == 1 and getattr(vectype, "_wp_generic_type_str_", None) == "vec_t":
+                # constructor from another vector
+                veclen = vectype._length_
+                vectype = vectype._wp_scalar_type_
+            elif not all(vectype == t for t in arg_types):
                 raise RuntimeError(
-                    f"All numeric arguments to vec() constructor should have the same type, expected {veclen} args of type {vectype}, received { ','.join(map(lambda x : str(x.type), args)) }"
+                    f"All numeric arguments to vec() constructor should have the same type, expected {veclen} arg_types of type {vectype}, received { ','.join(map(lambda t : str(t), arg_types)) }"
                 )
 
         # update the templates list, so we can generate vec<len, type>() correctly in codegen
@@ -629,9 +698,15 @@ def vector_constructor_func(args, kwds, templates):
     else:
         # construction of a predeclared type, e.g.: vec5d
         veclen, vectype = templates
-        if not all(vectype == a.type for a in args):
+        if len(arg_types) == 1 and getattr(arg_types[0], "_wp_generic_type_str_", None) == "vec_t":
+            # constructor from another vector
+            if arg_types[0]._length_ != veclen:
+                raise RuntimeError(
+                    f"Incompatible matrix sizes for casting copy constructor, {veclen} vs {arg_types[0]._length_}"
+                )
+        elif not all(vectype == t for t in arg_types):
             raise RuntimeError(
-                f"All numeric arguments to vec() constructor should have the same type, expected {veclen} args of type {vectype}, received { ','.join(map(lambda x : str(x.type), args)) }"
+                f"All numeric arguments to vec() constructor should have the same type, expected {veclen} arg_types of type {vectype}, received { ','.join(map(lambda t : str(t), arg_types)) }"
             )
 
     retvalue = vector(length=veclen, dtype=vectype)
@@ -640,9 +715,9 @@ def vector_constructor_func(args, kwds, templates):
 
 add_builtin(
     "vector",
-    input_types={"*args": Scalar, "length": int, "dtype": Scalar},
+    input_types={"*arg_types": Scalar, "length": int, "dtype": Scalar},
     variadic=True,
-    initializer_list_func=lambda args, _: len(args) > 4,
+    initializer_list_func=lambda arg_types, _: len(arg_types) > 4,
     value_func=vector_constructor_func,
     native_func="vec_t",
     doc="Construct a vector of with given length and dtype.",
@@ -651,8 +726,8 @@ add_builtin(
 )
 
 
-def matrix_constructor_func(args, kwds, templates):
-    if args is None:
+def matrix_constructor_func(arg_types, kwds, templates):
+    if arg_types is None:
         return matrix(shape=(Any, Any), dtype=Scalar)
 
     if len(templates) == 0:
@@ -660,7 +735,7 @@ def matrix_constructor_func(args, kwds, templates):
         if "shape" not in kwds:
             raise RuntimeError("shape keyword must be specified when calling matrix() function")
 
-        if len(args) == 0:
+        if len(arg_types) == 0:
             if "dtype" not in kwds:
                 raise RuntimeError("matrix() must have dtype as a keyword argument if it has no positional arguments")
 
@@ -671,9 +746,16 @@ def matrix_constructor_func(args, kwds, templates):
         else:
             # value initialization, e.g.: m = matrix(1.0, shape=(3,2))
             shape = kwds["shape"]
-            dtype = args[0].type
+            dtype = arg_types[0]
 
-            if len(args) > 1 and len(args) != shape[0] * shape[1]:
+            if len(arg_types) == 1 and getattr(dtype, "_wp_generic_type_str_", None) == "mat_t":
+                # constructor from another matrix
+                if arg_types[0]._shape_ != shape:
+                    raise RuntimeError(
+                        f"Incompatible matrix sizes for casting copy constructor, {shape} vs {arg_types[0]._shape_}"
+                    )
+                dtype = dtype._wp_scalar_type_
+            elif len(arg_types) > 1 and len(arg_types) != shape[0] * shape[1]:
                 raise RuntimeError(
                     "Wrong number of arguments for matrix() function, must initialize with either a scalar value, or m*n values"
                 )
@@ -687,68 +769,70 @@ def matrix_constructor_func(args, kwds, templates):
         shape = (templates[0], templates[1])
         dtype = templates[2]
 
-        if len(args) > 0:
-            # check scalar arg type matches declared type
-            if infer_scalar_type(args) != dtype:
-                raise RuntimeError("Wrong scalar type for mat {} constructor".format(",".join(map(str, templates))))
-
-            # check vector arg type matches declared type
-            types = [a.type for a in args]
-            if all(hasattr(a, "_wp_generic_type_str_") and a._wp_generic_type_str_ == "vec_t" for a in types):
-                cols = len(types)
-                if shape[1] != cols:
-                    raise RuntimeError(
-                        "Wrong number of vectors when attempting to construct a matrix with column vectors"
-                    )
-
-                if not all(a._length_ == shape[0] for a in types):
+        if len(arg_types) > 0:
+            if len(arg_types) == 1 and getattr(arg_types[0], "_wp_generic_type_str_", None) == "mat_t":
+                # constructor from another matrix with same dimension but possibly different type
+                if arg_types[0]._shape_ != shape:
                     raise RuntimeError(
-                        "Wrong vector row count when attempting to construct a matrix with column vectors"
+                        f"Incompatible matrix sizes for casting copy constructor, {shape} vs {arg_types[0]._shape_}"
                     )
-
             else:
-                # check that we either got 1 arg (scalar construction), or enough values for whole matrix
-                size = shape[0] * shape[1]
-                if len(args) > 1 and len(args) != size:
-                    raise RuntimeError(
-                        "Wrong number of scalars when attempting to construct a matrix from a list of components"
-                    )
+                # check scalar arg type matches declared type
+                if infer_scalar_type(arg_types) != dtype:
+                    raise RuntimeError("Wrong scalar type for mat {} constructor".format(",".join(map(str, templates))))
+
+                # check vector arg type matches declared type
+                if all(hasattr(a, "_wp_generic_type_str_") and a._wp_generic_type_str_ == "vec_t" for a in arg_types):
+                    cols = len(arg_types)
+                    if shape[1] != cols:
+                        raise RuntimeError(
+                            "Wrong number of vectors when attempting to construct a matrix with column vectors"
+                        )
+
+                    if not all(a._length_ == shape[0] for a in arg_types):
+                        raise RuntimeError(
+                            "Wrong vector row count when attempting to construct a matrix with column vectors"
+                        )
+                else:
+                    # check that we either got 1 arg (scalar construction), or enough values for whole matrix
+                    size = shape[0] * shape[1]
+                    if len(arg_types) > 1 and len(arg_types) != size:
+                        raise RuntimeError(
+                            "Wrong number of scalars when attempting to construct a matrix from a list of components"
+                        )
 
     return matrix(shape=shape, dtype=dtype)
 
 
 # only use initializer list if matrix size < 5x5, or for scalar construction
-def matrix_initlist_func(args, templates):
+def matrix_initlist_func(arg_types, templates):
     m, n, dtype = templates
-    if (
-        len(args) == 0
-        or len(args) == 1  # zero construction
+    return not (
+        len(arg_types) == 0
+        or len(arg_types) == 1  # zero construction
         or (m == n and n < 5)  # scalar construction  # value construction for small matrices
-    ):
-        return False
-    else:
-        return True
+    )
 
 
 add_builtin(
     "matrix",
-    input_types={"*args": Scalar, "shape": Tuple[int, int], "dtype": Scalar},
+    input_types={"*arg_types": Scalar, "shape": Tuple[int, int], "dtype": Scalar},
     variadic=True,
     initializer_list_func=matrix_initlist_func,
     value_func=matrix_constructor_func,
     native_func="mat_t",
-    doc="Construct a matrix, if positional args are not given then matrix will be zero-initialized.",
+    doc="Construct a matrix. If the positional ``arg_types`` are not given, then matrix will be zero-initialized.",
     group="Vector Math",
     export=False,
 )
 
 
 # identity:
-def matrix_identity_value_func(args, kwds, templates):
-    if args is None:
+def matrix_identity_value_func(arg_types, kwds, templates):
+    if arg_types is None:
         return matrix(shape=(Any, Any), dtype=Scalar)
 
-    if len(args):
+    if len(arg_types):
         raise RuntimeError("identity() function does not accept positional arguments")
 
     if "n" not in kwds:
@@ -779,7 +863,7 @@ add_builtin(
 )
 
 
-def matrix_transform_value_func(args, kwds, templates):
+def matrix_transform_value_func(arg_types, kwds, templates):
     if templates is None:
         return matrix(shape=(Any, Any), dtype=Float)
 
@@ -789,7 +873,7 @@ def matrix_transform_value_func(args, kwds, templates):
     m, n, dtype = templates
     if (m, n) != (4, 4):
         raise RuntimeError("Can only construct 4x4 matrices with position, rotation and scale")
-    if infer_scalar_type(args) != dtype:
+    if infer_scalar_type(arg_types) != dtype:
         raise RuntimeError("Wrong scalar type for mat<{}> constructor".format(",".join(map(str, templates))))
 
     return matrix(shape=(4, 4), dtype=dtype)
@@ -804,7 +888,8 @@ add_builtin(
     },
     value_func=matrix_transform_value_func,
     native_func="mat_t",
-    doc="""Construct a 4x4 transformation matrix that applies the transformations as Translation(pos)*Rotation(rot)*Scale(scale) when applied to column vectors, i.e.: y = (TRS)*x""",
+    doc="""Construct a 4x4 transformation matrix that applies the transformations as
+   Translation(pos)*Rotation(rot)*Scale(scale) when applied to column vectors, i.e.: y = (TRS)*x""",
     group="Vector Math",
     export=False,
 )
@@ -823,8 +908,8 @@ add_builtin(
     value_type=None,
     group="Vector Math",
     export=False,
-    doc="""Compute the SVD of a 3x3 matrix. The singular values are returned in sigma, 
-   while the left and right basis vectors are returned in U and V.""",
+    doc="""Compute the SVD of a 3x3 matrix ``A``. The singular values are returned in ``sigma``,
+   while the left and right basis vectors are returned in ``U`` and ``V``.""",
 )
 
 add_builtin(
@@ -837,7 +922,8 @@ add_builtin(
     value_type=None,
     group="Vector Math",
     export=False,
-    doc="""Compute the QR decomposition of a 3x3 matrix. The orthogonal matrix is returned in Q, while the upper triangular matrix is returned in R.""",
+    doc="""Compute the QR decomposition of a 3x3 matrix ``A``. The orthogonal matrix is returned in ``Q``,
+   while the upper triangular matrix is returned in ``R``.""",
 )
 
 add_builtin(
@@ -850,36 +936,53 @@ add_builtin(
     value_type=None,
     group="Vector Math",
     export=False,
-    doc="""Compute the eigendecomposition of a 3x3 matrix. The eigenvectors are returned as the columns of Q, while the corresponding eigenvalues are returned in d.""",
+    doc="""Compute the eigendecomposition of a 3x3 matrix ``A``. The eigenvectors are returned as the columns of ``Q``,
+   while the corresponding eigenvalues are returned in ``d``.""",
 )
 
 # ---------------------------------
 # Quaternion Math
 
 
-def quaternion_value_func(args, kwds, templates):
-    if args is None:
-        return quaternion(dtype=Scalar)
+def quaternion_value_func(arg_types, kwds, templates):
+    if arg_types is None:
+        return quaternion(dtype=Float)
 
-    # if constructing anonymous quat type then infer output type from arguments
     if len(templates) == 0:
-        dtype = infer_scalar_type(args)
+        if "dtype" in kwds:
+            # casting constructor
+            dtype = kwds["dtype"]
+        else:
+            # if constructing anonymous quat type then infer output type from arguments
+            dtype = infer_scalar_type(arg_types)
         templates.append(dtype)
     else:
-        # if constructing predeclared type then check args match expectation
-        if len(args) > 0 and infer_scalar_type(args) != templates[0]:
+        # if constructing predeclared type then check arg_types match expectation
+        if len(arg_types) > 0 and infer_scalar_type(arg_types) != templates[0]:
             raise RuntimeError("Wrong scalar type for quat {} constructor".format(",".join(map(str, templates))))
 
     return quaternion(dtype=templates[0])
 
 
+def quat_cast_value_func(arg_types, kwds, templates):
+    if arg_types is None:
+        raise RuntimeError("Missing quaternion argument.")
+    if "dtype" not in kwds:
+        raise RuntimeError("Missing 'dtype' kwd.")
+
+    dtype = kwds["dtype"]
+    templates.append(dtype)
+
+    return quaternion(dtype=dtype)
+
+
 add_builtin(
     "quaternion",
     input_types={},
     value_func=quaternion_value_func,
     native_func="quat_t",
     group="Quaternion Math",
-    doc="""Construct a zero-initialized quaternion, quaternions are laid out as
+    doc="""Construct a zero-initialized quaternion. Quaternions are laid out as
    [ix, iy, iz, r], where ix, iy, iz are the imaginary part, and r the real part.""",
     export=False,
 )
@@ -889,7 +992,7 @@ add_builtin(
     value_func=quaternion_value_func,
     native_func="quat_t",
     group="Quaternion Math",
-    doc="Create a quaternion using the supplied components (type inferred from component type)",
+    doc="Create a quaternion using the supplied components (type inferred from component type).",
     export=False,
 )
 add_builtin(
@@ -898,14 +1001,23 @@ add_builtin(
     value_func=quaternion_value_func,
     native_func="quat_t",
     group="Quaternion Math",
-    doc="Create a quaternion using the supplied vector/scalar (type inferred from scalar type)",
+    doc="Create a quaternion using the supplied vector/scalar (type inferred from scalar type).",
+    export=False,
+)
+add_builtin(
+    "quaternion",
+    input_types={"q": quaternion(dtype=Float)},
+    value_func=quat_cast_value_func,
+    native_func="quat_t",
+    group="Quaternion Math",
+    doc="Construct a quaternion of type dtype from another quaternion of a different dtype.",
     export=False,
 )
 
 
-def quat_identity_value_func(args, kwds, templates):
-    # if args is None then we are in 'export' mode
-    if args is None:
+def quat_identity_value_func(arg_types, kwds, templates):
+    # if arg_types is None then we are in 'export' mode
+    if arg_types is None:
         return quatf
 
     if "dtype" not in kwds:
@@ -931,7 +1043,7 @@ add_builtin(
 add_builtin(
     "quat_from_axis_angle",
     input_types={"axis": vector(length=3, dtype=Float), "angle": Float},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Construct a quaternion representing a rotation of angle radians around the given axis.",
 )
@@ -945,49 +1057,50 @@ add_builtin(
 add_builtin(
     "quat_from_matrix",
     input_types={"m": matrix(shape=(3, 3), dtype=Float)},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Construct a quaternion from a 3x3 matrix.",
 )
 add_builtin(
     "quat_rpy",
     input_types={"roll": Float, "pitch": Float, "yaw": Float},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Construct a quaternion representing a combined roll (z), pitch (x), yaw rotations (y) in radians.",
 )
 add_builtin(
     "quat_inverse",
     input_types={"q": quaternion(dtype=Float)},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Compute quaternion conjugate.",
 )
 add_builtin(
     "quat_rotate",
     input_types={"q": quaternion(dtype=Float), "p": vector(length=3, dtype=Float)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Rotate a vector by a quaternion.",
 )
 add_builtin(
     "quat_rotate_inv",
     input_types={"q": quaternion(dtype=Float), "p": vector(length=3, dtype=Float)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
-    doc="Rotate a vector the inverse of a quaternion.",
+    doc="Rotate a vector by the inverse of a quaternion.",
 )
 add_builtin(
     "quat_slerp",
     input_types={"q0": quaternion(dtype=Float), "q1": quaternion(dtype=Float), "t": Float},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Linearly interpolate between two quaternions.",
+    require_original_output_arg=True,
 )
 add_builtin(
     "quat_to_matrix",
     input_types={"q": quaternion(dtype=Float)},
-    value_func=lambda args, kwds, _: matrix(shape=(3, 3), dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: matrix(shape=(3, 3), dtype=infer_scalar_type(arg_types)),
     group="Quaternion Math",
     doc="Convert a quaternion to a 3x3 rotation matrix.",
 )
@@ -1003,19 +1116,19 @@ add_builtin(
 # Transformations
 
 
-def transform_constructor_value_func(args, kwds, templates):
+def transform_constructor_value_func(arg_types, kwds, templates):
     if templates is None:
         return transformation(dtype=Scalar)
 
     if len(templates) == 0:
         # if constructing anonymous transform type then infer output type from arguments
-        dtype = infer_scalar_type(args)
+        dtype = infer_scalar_type(arg_types)
         templates.append(dtype)
     else:
-        # if constructing predeclared type then check args match expectation
-        if infer_scalar_type(args) != templates[0]:
+        # if constructing predeclared type then check arg_types match expectation
+        if infer_scalar_type(arg_types) != templates[0]:
             raise RuntimeError(
-                f"Wrong scalar type for transform constructor expected {templates[0]}, got {','.join(map(lambda x : str(x.type), args))}"
+                f"Wrong scalar type for transform constructor expected {templates[0]}, got {','.join(map(lambda t : str(t), arg_types))}"
             )
 
     return transformation(dtype=templates[0])
@@ -1027,13 +1140,13 @@ add_builtin(
     value_func=transform_constructor_value_func,
     native_func="transform_t",
     group="Transformations",
-    doc="Construct a rigid body transformation with translation part p and rotation q.",
+    doc="Construct a rigid-body transformation with translation part ``p`` and rotation ``q``.",
     export=False,
 )
 
 
-def transform_identity_value_func(args, kwds, templates):
-    if args is None:
+def transform_identity_value_func(arg_types, kwds, templates):
+    if arg_types is None:
         return transformf
 
     if "dtype" not in kwds:
@@ -1059,68 +1172,72 @@ add_builtin(
 add_builtin(
     "transform_get_translation",
     input_types={"t": transformation(dtype=Float)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Transformations",
-    doc="Return the translational part of a transform.",
+    doc="Return the translational part of a transform ``t``.",
 )
 add_builtin(
     "transform_get_rotation",
     input_types={"t": transformation(dtype=Float)},
-    value_func=lambda args, kwds, _: quaternion(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: quaternion(dtype=infer_scalar_type(arg_types)),
     group="Transformations",
-    doc="Return the rotational part of a transform.",
+    doc="Return the rotational part of a transform ``t``.",
 )
 add_builtin(
     "transform_multiply",
     input_types={"a": transformation(dtype=Float), "b": transformation(dtype=Float)},
-    value_func=lambda args, kwds, _: transformation(dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: transformation(dtype=infer_scalar_type(arg_types)),
     group="Transformations",
     doc="Multiply two rigid body transformations together.",
 )
 add_builtin(
     "transform_point",
     input_types={"t": transformation(dtype=Scalar), "p": vector(length=3, dtype=Scalar)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Transformations",
-    doc="Apply the transform to a point p treating the homogenous coordinate as w=1 (translation and rotation).",
+    doc="Apply the transform to a point ``p`` treating the homogeneous coordinate as w=1 (translation and rotation).",
 )
 add_builtin(
     "transform_point",
     input_types={"m": matrix(shape=(4, 4), dtype=Scalar), "p": vector(length=3, dtype=Scalar)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Vector Math",
-    doc="""Apply the transform to a point ``p`` treating the homogenous coordinate as w=1. The transformation is applied treating ``p`` as a column vector, e.g.: ``y = M*p``
-   note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = p^T*M^T``. If the transform is coming from a library that uses row-vectors
-   then users should transpose the transformation matrix before calling this method.""",
+    doc="""Apply the transform to a point ``p`` treating the homogeneous coordinate as w=1.
+   The transformation is applied treating ``p`` as a column vector, e.g.: ``y = M*p``.
+   Note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = p^T*M^T``.
+   If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
+   matrix before calling this method.""",
 )
 add_builtin(
     "transform_vector",
     input_types={"t": transformation(dtype=Scalar), "v": vector(length=3, dtype=Scalar)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Transformations",
-    doc="Apply the transform to a vector v treating the homogenous coordinate as w=0 (rotation only).",
+    doc="Apply the transform to a vector ``v`` treating the homogeneous coordinate as w=0 (rotation only).",
 )
 add_builtin(
     "transform_vector",
     input_types={"m": matrix(shape=(4, 4), dtype=Scalar), "v": vector(length=3, dtype=Scalar)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=infer_scalar_type(arg_types)),
     group="Vector Math",
-    doc="""Apply the transform to a vector ``v`` treating the homogenous coordinate as w=0. The transformation is applied treating ``v`` as a column vector, e.g.: ``y = M*v``
-   note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = v^T*M^T``. If the transform is coming from a library that uses row-vectors
-   then users should transpose the transformation matrix before calling this method.""",
+    doc="""Apply the transform to a vector ``v`` treating the homogeneous coordinate as w=0.
+   The transformation is applied treating ``v`` as a column vector, e.g.: ``y = M*v``
+   note this is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = v^T*M^T``.
+   If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
+   matrix before calling this method.""",
 )
 add_builtin(
     "transform_inverse",
     input_types={"t": transformation(dtype=Float)},
     value_func=sametype_value_func(transformation(dtype=Float)),
     group="Transformations",
-    doc="Compute the inverse of the transform.",
+    doc="Compute the inverse of the transformation ``t``.",
 )
 # ---------------------------------
 # Spatial Math
 
 
-def spatial_vector_constructor_value_func(args, kwds, templates):
+def spatial_vector_constructor_value_func(arg_types, kwds, templates):
     if templates is None:
         return spatial_vector(dtype=Float)
 
@@ -1128,7 +1245,7 @@ def spatial_vector_constructor_value_func(args, kwds, templates):
         raise RuntimeError("Cannot use a generic type name in a kernel")
 
     vectype = templates[1]
-    if len(args) and infer_scalar_type(args) != vectype:
+    if len(arg_types) and infer_scalar_type(arg_types) != vectype:
         raise RuntimeError("Wrong scalar type for spatial_vector<{}> constructor".format(",".join(map(str, templates))))
 
     return vector(length=6, dtype=vectype)
@@ -1140,7 +1257,7 @@ add_builtin(
     value_func=spatial_vector_constructor_value_func,
     native_func="vec_t",
     group="Spatial Math",
-    doc="Construct a 6d screw vector from two 3d vectors.",
+    doc="Construct a 6D screw vector from two 3D vectors.",
     export=False,
 )
 
@@ -1148,7 +1265,7 @@ add_builtin(
 add_builtin(
     "spatial_adjoint",
     input_types={"r": matrix(shape=(3, 3), dtype=Float), "s": matrix(shape=(3, 3), dtype=Float)},
-    value_func=lambda args, kwds, _: matrix(shape=(6, 6), dtype=infer_scalar_type(args)),
+    value_func=lambda arg_types, kwds, _: matrix(shape=(6, 6), dtype=infer_scalar_type(arg_types)),
     group="Spatial Math",
     doc="Construct a 6x6 spatial inertial matrix from two 3x3 diagonal blocks.",
     export=False,
@@ -1158,36 +1275,36 @@ add_builtin(
     input_types={"a": vector(length=6, dtype=Float), "b": vector(length=6, dtype=Float)},
     value_func=sametype_scalar_value_func,
     group="Spatial Math",
-    doc="Compute the dot product of two 6d screw vectors.",
+    doc="Compute the dot product of two 6D screw vectors.",
 )
 add_builtin(
     "spatial_cross",
     input_types={"a": vector(length=6, dtype=Float), "b": vector(length=6, dtype=Float)},
     value_func=sametype_value_func(vector(length=6, dtype=Float)),
     group="Spatial Math",
-    doc="Compute the cross-product of two 6d screw vectors.",
+    doc="Compute the cross product of two 6D screw vectors.",
 )
 add_builtin(
     "spatial_cross_dual",
     input_types={"a": vector(length=6, dtype=Float), "b": vector(length=6, dtype=Float)},
     value_func=sametype_value_func(vector(length=6, dtype=Float)),
     group="Spatial Math",
-    doc="Compute the dual cross-product of two 6d screw vectors.",
+    doc="Compute the dual cross product of two 6D screw vectors.",
 )
 
 add_builtin(
     "spatial_top",
     input_types={"a": vector(length=6, dtype=Float)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=args[0].type._wp_scalar_type_),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=arg_types[0]._wp_scalar_type_),
     group="Spatial Math",
-    doc="Return the top (first) part of a 6d screw vector.",
+    doc="Return the top (first) part of a 6D screw vector.",
 )
 add_builtin(
     "spatial_bottom",
     input_types={"a": vector(length=6, dtype=Float)},
-    value_func=lambda args, kwds, _: vector(length=3, dtype=args[0].type._wp_scalar_type_),
+    value_func=lambda arg_types, kwds, _: vector(length=3, dtype=arg_types[0]._wp_scalar_type_),
     group="Spatial Math",
-    doc="Return the bottom (second) part of a 6d screw vector.",
+    doc="Return the bottom (second) part of a 6D screw vector.",
 )
 
 add_builtin(
@@ -1341,16 +1458,18 @@ add_builtin(
     },
     value_type=None,
     skip_replay=True,
-    doc="""Evaluate a multi-layer perceptron (MLP) layer in the form: ``out = act(weights*x + bias)``. 
+    doc="""Evaluate a multi-layer perceptron (MLP) layer in the form: ``out = act(weights*x + bias)``.
 
    :param weights: A layer's network weights with dimensions ``(m, n)``.
    :param bias: An array with dimensions ``(n)``.
    :param activation: A ``wp.func`` function that takes a single scalar float as input and returns a scalar float as output
-   :param index: The batch item to process, typically each thread will process 1 item in the batch, in this case index should be ``wp.tid()``
+   :param index: The batch item to process, typically each thread will process one item in the batch, in which case
+                 index should be ``wp.tid()``
    :param x: The feature matrix with dimensions ``(n, b)``
    :param out: The network output with dimensions ``(m, b)``
 
-   :note: Feature and output matrices are transposed compared to some other frameworks such as PyTorch. All matrices are assumed to be stored in flattened row-major memory layout (NumPy default).""",
+   :note: Feature and output matrices are transposed compared to some other frameworks such as PyTorch.
+          All matrices are assumed to be stored in flattened row-major memory layout (NumPy default).""",
     group="Utility",
 )
 
@@ -1363,12 +1482,12 @@ add_builtin(
     input_types={"id": uint64, "lower": vec3, "upper": vec3},
     value_type=bvh_query_t,
     group="Geometry",
-    doc="""Construct an axis-aligned bounding box query against a bvh object. This query can be used to iterate over all bounds
-   inside a bvh. Returns an object that is used to track state during bvh traversal.
-    
-   :param id: The bvh identifier
-   :param lower: The lower bound of the bounding box in bvh space
-   :param upper: The upper bound of the bounding box in bvh space""",
+    doc="""Construct an axis-aligned bounding box query against a BVH object. This query can be used to iterate over all bounds
+   inside a BVH.
+
+   :param id: The BVH identifier
+   :param lower: The lower bound of the bounding box in BVH space
+   :param upper: The upper bound of the bounding box in BVH space""",
 )
 
 add_builtin(
@@ -1376,21 +1495,21 @@ add_builtin(
     input_types={"id": uint64, "start": vec3, "dir": vec3},
     value_type=bvh_query_t,
     group="Geometry",
-    doc="""Construct a ray query against a bvh object. This query can be used to iterate over all bounds
-   that intersect the ray. Returns an object that is used to track state during bvh traversal.
-    
-   :param id: The bvh identifier
-   :param start: The start of the ray in bvh space
-   :param dir: The direction of the ray in bvh space""",
+    doc="""Construct a ray query against a BVH object. This query can be used to iterate over all bounds
+   that intersect the ray.
+
+   :param id: The BVH identifier
+   :param start: The start of the ray in BVH space
+   :param dir: The direction of the ray in BVH space""",
 )
 
 add_builtin(
     "bvh_query_next",
     input_types={"query": bvh_query_t, "index": int},
-    value_type=bool,
+    value_type=builtins.bool,
     group="Geometry",
-    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.""",
+    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.""",
 )
 
 add_builtin(
@@ -1404,17 +1523,256 @@ add_builtin(
         "bary_u": float,
         "bary_v": float,
     },
-    value_type=bool,
+    value_type=builtins.bool,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+
+   Identifies the sign of the distance using additional ray-casts to determine if the point is inside or outside.
+   This method is relatively robust, but does increase computational cost.
+   See below for additional sign determination methods.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query
+   :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
+                  Note that mesh must be watertight for this to be robust
+   :param face: Returns the index of the closest face
+   :param bary_u: Returns the barycentric u coordinate of the closest point
+   :param bary_v: Returns the barycentric v coordinate of the closest point""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_point",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+    },
+    value_type=mesh_query_point_t,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
+
+   Identifies the sign of the distance using additional ray-casts to determine if the point is inside or outside.
+   This method is relatively robust, but does increase computational cost.
+   See below for additional sign determination methods.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query""",
+    require_original_output_arg=True,
+)
+
+add_builtin(
+    "mesh_query_point_no_sign",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+        "face": int,
+        "bary_u": float,
+        "bary_v": float,
+    },
+    value_type=builtins.bool,
     group="Geometry",
-    doc="""Computes the closest point on the mesh with identifier `id` to the given point in space. Returns ``True`` if a point < ``max_dist`` is found.
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+
+   This method does not compute the sign of the point (inside/outside) which makes it faster than other point query methods.
 
    :param id: The mesh identifier
    :param point: The point in space to query
    :param max_dist: Mesh faces above this distance will not be considered by the query
-   :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise. Note that mesh must be watertight for this to be robust
    :param face: Returns the index of the closest face
    :param bary_u: Returns the barycentric u coordinate of the closest point
    :param bary_v: Returns the barycentric v coordinate of the closest point""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_point_no_sign",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+    },
+    value_type=mesh_query_point_t,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
+
+   This method does not compute the sign of the point (inside/outside) which makes it faster than other point query methods.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query""",
+    require_original_output_arg=True,
+)
+
+add_builtin(
+    "mesh_query_furthest_point_no_sign",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "min_dist": float,
+        "face": int,
+        "bary_u": float,
+        "bary_v": float,
+    },
+    value_type=builtins.bool,
+    group="Geometry",
+    doc="""Computes the furthest point on the mesh with identifier `id` to the given point in space. Returns ``True`` if a point > ``min_dist`` is found.
+
+   This method does not compute the sign of the point (inside/outside).
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param min_dist: Mesh faces below this distance will not be considered by the query
+   :param face: Returns the index of the furthest face
+   :param bary_u: Returns the barycentric u coordinate of the furthest point
+   :param bary_v: Returns the barycentric v coordinate of the furthest point""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_furthest_point_no_sign",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "min_dist": float,
+    },
+    value_type=mesh_query_point_t,
+    group="Geometry",
+    doc="""Computes the furthest point on the mesh with identifier `id` to the given point in space.
+
+   This method does not compute the sign of the point (inside/outside).
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param min_dist: Mesh faces below this distance will not be considered by the query""",
+    require_original_output_arg=True,
+)
+
+add_builtin(
+    "mesh_query_point_sign_normal",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+        "inside": float,
+        "face": int,
+        "bary_u": float,
+        "bary_v": float,
+        "epsilon": float,
+    },
+    defaults={"epsilon": 1.0e-3},
+    value_type=builtins.bool,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space. Returns ``True`` if a point < ``max_dist`` is found.
+
+   Identifies the sign of the distance (inside/outside) using the angle-weighted pseudo normal.
+   This approach to sign determination is robust for well conditioned meshes that are watertight and non-self intersecting.
+   It is also comparatively fast to compute.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query
+   :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
+                  Note that mesh must be watertight for this to be robust
+   :param face: Returns the index of the closest face
+   :param bary_u: Returns the barycentric u coordinate of the closest point
+   :param bary_v: Returns the barycentric v coordinate of the closest point
+   :param epsilon: Epsilon treating distance values as equal, when locating the minimum distance vertex/face/edge, as a
+                   fraction of the average edge length, also for treating closest point as being on edge/vertex default 1e-3""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_point_sign_normal",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+        "epsilon": float,
+    },
+    defaults={"epsilon": 1.0e-3},
+    value_type=mesh_query_point_t,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
+
+   Identifies the sign of the distance (inside/outside) using the angle-weighted pseudo normal.
+   This approach to sign determination is robust for well conditioned meshes that are watertight and non-self intersecting.
+   It is also comparatively fast to compute.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query
+   :param epsilon: Epsilon treating distance values as equal, when locating the minimum distance vertex/face/edge, as a
+                   fraction of the average edge length, also for treating closest point as being on edge/vertex default 1e-3""",
+    require_original_output_arg=True,
+)
+
+add_builtin(
+    "mesh_query_point_sign_winding_number",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+        "inside": float,
+        "face": int,
+        "bary_u": float,
+        "bary_v": float,
+        "accuracy": float,
+        "threshold": float,
+    },
+    defaults={"accuracy": 2.0, "threshold": 0.5},
+    value_type=builtins.bool,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given point in space. Returns ``True`` if a point < ``max_dist`` is found.
+
+   Identifies the sign using the winding number of the mesh relative to the query point. This method of sign determination is robust for poorly conditioned meshes
+   and provides a smooth approximation to sign even when the mesh is not watertight. This method is the most robust and accurate of the sign determination meshes
+   but also the most expensive.
+
+   .. note:: The :class:`Mesh` object must be constructed with ``support_winding_number=True`` for this method to return correct results.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query
+   :param inside: Returns a value < 0 if query point is inside the mesh, >=0 otherwise.
+                  Note that mesh must be watertight for this to be robust
+   :param face: Returns the index of the closest face
+   :param bary_u: Returns the barycentric u coordinate of the closest point
+   :param bary_v: Returns the barycentric v coordinate of the closest point
+   :param accuracy: Accuracy for computing the winding number with fast winding number method utilizing second-order dipole approximation, default 2.0
+   :param threshold: The threshold of the winding number to be considered inside, default 0.5""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_point_sign_winding_number",
+    input_types={
+        "id": uint64,
+        "point": vec3,
+        "max_dist": float,
+        "accuracy": float,
+        "threshold": float,
+    },
+    defaults={"accuracy": 2.0, "threshold": 0.5},
+    value_type=mesh_query_point_t,
+    group="Geometry",
+    doc="""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given point in space.
+
+   Identifies the sign using the winding number of the mesh relative to the query point. This method of sign determination is robust for poorly conditioned meshes
+   and provides a smooth approximation to sign even when the mesh is not watertight. This method is the most robust and accurate of the sign determination meshes
+   but also the most expensive.
+
+   .. note:: The :class:`Mesh` object must be constructed with ``support_winding_number=True`` for this method to return correct results.
+
+   :param id: The mesh identifier
+   :param point: The point in space to query
+   :param max_dist: Mesh faces above this distance will not be considered by the query
+   :param accuracy: Accuracy for computing the winding number with fast winding number method utilizing second-order dipole approximation, default 2.0
+   :param threshold: The threshold of the winding number to be considered inside, default 0.5""",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -1431,9 +1789,9 @@ add_builtin(
         "normal": vec3,
         "face": int,
     },
-    value_type=bool,
+    value_type=builtins.bool,
     group="Geometry",
-    doc="""Computes the closest ray hit on the mesh with identifier `id`, returns ``True`` if a point < ``max_t`` is found.
+    doc="""Computes the closest ray hit on the :class:`Mesh` with identifier ``id``, returns ``True`` if a hit < ``max_t`` is found.
 
    :param id: The mesh identifier
    :param start: The start point of the ray
@@ -1442,9 +1800,29 @@ add_builtin(
    :param t: Returns the distance of the closest hit along the ray
    :param bary_u: Returns the barycentric u coordinate of the closest hit
    :param bary_v: Returns the barycentric v coordinate of the closest hit
-   :param sign: Returns a value > 0 if the hit ray hit front of the face, returns < 0 otherwise
+   :param sign: Returns a value > 0 if the ray hit in front of the face, returns < 0 otherwise
    :param normal: Returns the face normal
    :param face: Returns the index of the hit face""",
+    hidden=True,
+)
+
+add_builtin(
+    "mesh_query_ray",
+    input_types={
+        "id": uint64,
+        "start": vec3,
+        "dir": vec3,
+        "max_t": float,
+    },
+    value_type=mesh_query_ray_t,
+    group="Geometry",
+    doc="""Computes the closest ray hit on the :class:`Mesh` with identifier ``id``.
+
+   :param id: The mesh identifier
+   :param start: The start point of the ray
+   :param dir: The ray direction (should be normalized)
+   :param max_t: The maximum distance along the ray to check for intersections""",
+    require_original_output_arg=True,
 )
 
 add_builtin(
@@ -1452,9 +1830,9 @@ add_builtin(
     input_types={"id": uint64, "lower": vec3, "upper": vec3},
     value_type=mesh_query_aabb_t,
     group="Geometry",
-    doc="""Construct an axis-aligned bounding box query against a mesh object. This query can be used to iterate over all triangles
-   inside a volume. Returns an object that is used to track state during mesh traversal.
-    
+    doc="""Construct an axis-aligned bounding box query against a :class:`Mesh`.
+   This query can be used to iterate over all triangles inside a volume.
+
    :param id: The mesh identifier
    :param lower: The lower bound of the bounding box in mesh space
    :param upper: The upper bound of the bounding box in mesh space""",
@@ -1463,10 +1841,10 @@ add_builtin(
 add_builtin(
     "mesh_query_aabb_next",
     input_types={"query": mesh_query_aabb_t, "index": int},
-    value_type=bool,
+    value_type=builtins.bool,
     group="Geometry",
-    doc="""Move to the next triangle overlapping the query bounding box. The index of the current face is stored in ``index``, returns ``False``
-   if there are no more overlapping triangles.""",
+    doc="""Move to the next triangle overlapping the query bounding box.
+   The index of the current face is stored in ``index``, returns ``False`` if there are no more overlapping triangles.""",
 )
 
 add_builtin(
@@ -1474,7 +1852,7 @@ add_builtin(
     input_types={"id": uint64, "face": int, "bary_u": float, "bary_v": float},
     value_type=vec3,
     group="Geometry",
-    doc="""Evaluates the position on the mesh given a face index, and barycentric coordinates.""",
+    doc="""Evaluates the position on the :class:`Mesh` given a face index and barycentric coordinates.""",
 )
 
 add_builtin(
@@ -1482,7 +1860,7 @@ add_builtin(
     input_types={"id": uint64, "face": int, "bary_u": float, "bary_v": float},
     value_type=vec3,
     group="Geometry",
-    doc="""Evaluates the velocity on the mesh given a face index, and barycentric coordinates.""",
+    doc="""Evaluates the velocity on the :class:`Mesh` given a face index and barycentric coordinates.""",
 )
 
 add_builtin(
@@ -1490,14 +1868,14 @@ add_builtin(
     input_types={"id": uint64, "point": vec3, "max_dist": float},
     value_type=hash_grid_query_t,
     group="Geometry",
-    doc="""Construct a point query against a hash grid. This query can be used to iterate over all neighboring points withing a 
-   fixed radius from the query point. Returns an object that is used to track state during neighbor traversal.""",
+    doc="Construct a point query against a :class:`HashGrid`. This query can be used to iterate over all neighboring points "
+    "within a fixed radius from the query point.",
 )
 
 add_builtin(
     "hash_grid_query_next",
     input_types={"query": hash_grid_query_t, "index": int},
-    value_type=bool,
+    value_type=builtins.bool,
     group="Geometry",
     doc="""Move to the next point in the hash grid query. The index of the current neighbor is stored in ``index``, returns ``False``
    if there are no more neighbors.""",
@@ -1508,8 +1886,10 @@ add_builtin(
     input_types={"id": uint64, "index": int},
     value_type=int,
     group="Geometry",
-    doc="""Return the index of a point in the grid, this can be used to re-order threads such that grid 
-   traversal occurs in a spatially coherent order.""",
+    doc="""Return the index of a point in the :class:`HashGrid`. This can be used to reorder threads such that grid
+   traversal occurs in a spatially coherent order.
+
+   Returns -1 if the :class:`HashGrid` has not been reserved.""",
 )
 
 add_builtin(
@@ -1608,7 +1988,17 @@ add_builtin(
     input_types={"id": uint64, "uvw": vec3, "sampling_mode": int},
     value_type=float,
     group="Volumes",
-    doc="""Sample the volume given by ``id`` at the volume local-space point ``uvw``. Interpolation should be ``wp.Volume.CLOSEST``, or ``wp.Volume.LINEAR.``""",
+    doc="""Sample the volume given by ``id`` at the volume local-space point ``uvw``.
+   Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
+)
+
+add_builtin(
+    "volume_sample_grad_f",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "grad": vec3},
+    value_type=float,
+    group="Volumes",
+    doc="""Sample the volume and its gradient given by ``id`` at the volume local-space point ``uvw``. 
+   Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
 )
 
 add_builtin(
@@ -1616,14 +2006,15 @@ add_builtin(
     input_types={"id": uint64, "i": int, "j": int, "k": int},
     value_type=float,
     group="Volumes",
-    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""",
+    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""",
 )
 
 add_builtin(
     "volume_store_f",
     input_types={"id": uint64, "i": int, "j": int, "k": int, "value": float},
     group="Volumes",
-    doc="""Store the value at voxel with coordinates ``i``, ``j``, ``k``.""",
+    doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
 )
 
 add_builtin(
@@ -1631,7 +2022,8 @@ add_builtin(
     input_types={"id": uint64, "uvw": vec3, "sampling_mode": int},
     value_type=vec3,
     group="Volumes",
-    doc="""Sample the vector volume given by ``id`` at the volume local-space point ``uvw``. Interpolation should be ``wp.Volume.CLOSEST``, or ``wp.Volume.LINEAR.``""",
+    doc="""Sample the vector volume given by ``id`` at the volume local-space point ``uvw``.
+   Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
 )
 
 add_builtin(
@@ -1639,14 +2031,15 @@ add_builtin(
     input_types={"id": uint64, "i": int, "j": int, "k": int},
     value_type=vec3,
     group="Volumes",
-    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""",
+    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.""",
 )
 
 add_builtin(
     "volume_store_v",
     input_types={"id": uint64, "i": int, "j": int, "k": int, "value": vec3},
     group="Volumes",
-    doc="""Store the value at voxel with coordinates ``i``, ``j``, ``k``.""",
+    doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
 )
 
 add_builtin(
@@ -1654,7 +2047,7 @@ add_builtin(
     input_types={"id": uint64, "uvw": vec3},
     value_type=int,
     group="Volumes",
-    doc="""Sample the int32 volume given by ``id`` at the volume local-space point ``uvw``. """,
+    doc="""Sample the :class:`int32` volume given by ``id`` at the volume local-space point ``uvw``. """,
 )
 
 add_builtin(
@@ -1662,14 +2055,15 @@ add_builtin(
     input_types={"id": uint64, "i": int, "j": int, "k": int},
     value_type=int,
     group="Volumes",
-    doc="""Returns the int32 value of voxel with coordinates ``i``, ``j``, ``k``, if the voxel at this index does not exist this function returns the background value""",
+    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.""",
 )
 
 add_builtin(
     "volume_store_i",
     input_types={"id": uint64, "i": int, "j": int, "k": int, "value": int},
     group="Volumes",
-    doc="""Store the value at voxel with coordinates ``i``, ``j``, ``k``.""",
+    doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
 )
 
 add_builtin(
@@ -1677,28 +2071,28 @@ add_builtin(
     input_types={"id": uint64, "uvw": vec3},
     value_type=vec3,
     group="Volumes",
-    doc="""Transform a point defined in volume index space to world space given the volume's intrinsic affine transformation.""",
+    doc="""Transform a point ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation.""",
 )
 add_builtin(
     "volume_world_to_index",
     input_types={"id": uint64, "xyz": vec3},
     value_type=vec3,
     group="Volumes",
-    doc="""Transform a point defined in volume world space to the volume's index space, given the volume's intrinsic affine transformation.""",
+    doc="""Transform a point ``xyz`` defined in volume world space to the volume's index space given the volume's intrinsic affine transformation.""",
 )
 add_builtin(
     "volume_index_to_world_dir",
     input_types={"id": uint64, "uvw": vec3},
     value_type=vec3,
     group="Volumes",
-    doc="""Transform a direction defined in volume index space to world space given the volume's intrinsic affine transformation.""",
+    doc="""Transform a direction ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation.""",
 )
 add_builtin(
     "volume_world_to_index_dir",
     input_types={"id": uint64, "xyz": vec3},
     value_type=vec3,
     group="Volumes",
-    doc="""Transform a direction defined in volume world space to the volume's index space, given the volume's intrinsic affine transformation.""",
+    doc="""Transform a direction ``xyz`` defined in volume world space to the volume's index space given the volume's intrinsic affine transformation.""",
 )
 
 
@@ -1718,7 +2112,7 @@ add_builtin(
     input_types={"seed": int, "offset": int},
     value_type=uint32,
     group="Random",
-    doc="""Initialize a new random number generator given a user-defined seed and an offset. 
+    doc="""Initialize a new random number generator given a user-defined seed and an offset.
    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)``""",
 )
@@ -1728,31 +2122,31 @@ add_builtin(
     input_types={"state": uint32},
     value_type=int,
     group="Random",
-    doc="Return a random integer between [0, 2^32)",
+    doc="Return a random integer in the range [0, 2^32).",
 )
 add_builtin(
     "randi",
     input_types={"state": uint32, "min": int, "max": int},
     value_type=int,
     group="Random",
-    doc="Return a random integer between [min, max)",
+    doc="Return a random integer between [min, max).",
 )
 add_builtin(
     "randf",
     input_types={"state": uint32},
     value_type=float,
     group="Random",
-    doc="Return a random float between [0.0, 1.0)",
+    doc="Return a random float between [0.0, 1.0).",
 )
 add_builtin(
     "randf",
     input_types={"state": uint32, "min": float, "max": float},
     value_type=float,
     group="Random",
-    doc="Return a random float between [min, max)",
+    doc="Return a random float between [min, max).",
 )
 add_builtin(
-    "randn", input_types={"state": uint32}, value_type=float, group="Random", doc="Sample a normal distribution"
+    "randn", input_types={"state": uint32}, value_type=float, group="Random", doc="Sample a normal distribution."
 )
 
 add_builtin(
@@ -1760,70 +2154,70 @@ add_builtin(
     input_types={"state": uint32, "cdf": array(dtype=float)},
     value_type=int,
     group="Random",
-    doc="Inverse transform sample a cumulative distribution function",
+    doc="Inverse-transform sample a cumulative distribution function.",
 )
 add_builtin(
     "sample_triangle",
     input_types={"state": uint32},
     value_type=vec2,
     group="Random",
-    doc="Uniformly sample a triangle. Returns sample barycentric coordinates",
+    doc="Uniformly sample a triangle. Returns sample barycentric coordinates.",
 )
 add_builtin(
     "sample_unit_ring",
     input_types={"state": uint32},
     value_type=vec2,
     group="Random",
-    doc="Uniformly sample a ring in the xy plane",
+    doc="Uniformly sample a ring in the xy plane.",
 )
 add_builtin(
     "sample_unit_disk",
     input_types={"state": uint32},
     value_type=vec2,
     group="Random",
-    doc="Uniformly sample a disk in the xy plane",
+    doc="Uniformly sample a disk in the xy plane.",
 )
 add_builtin(
     "sample_unit_sphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
     group="Random",
-    doc="Uniformly sample a unit sphere surface",
+    doc="Uniformly sample a unit sphere surface.",
 )
 add_builtin(
     "sample_unit_sphere",
     input_types={"state": uint32},
     value_type=vec3,
     group="Random",
-    doc="Uniformly sample a unit sphere",
+    doc="Uniformly sample a unit sphere.",
 )
 add_builtin(
     "sample_unit_hemisphere_surface",
     input_types={"state": uint32},
     value_type=vec3,
     group="Random",
-    doc="Uniformly sample a unit hemisphere surface",
+    doc="Uniformly sample a unit hemisphere surface.",
 )
 add_builtin(
     "sample_unit_hemisphere",
     input_types={"state": uint32},
     value_type=vec3,
     group="Random",
-    doc="Uniformly sample a unit hemisphere",
+    doc="Uniformly sample a unit hemisphere.",
 )
 add_builtin(
     "sample_unit_square",
     input_types={"state": uint32},
     value_type=vec2,
     group="Random",
-    doc="Uniformly sample a unit square",
+    doc="Uniformly sample a unit square.",
 )
 add_builtin(
     "sample_unit_cube",
     input_types={"state": uint32},
     value_type=vec3,
     group="Random",
-    doc="Uniformly sample a unit cube",
+    doc="Uniformly sample a unit cube.",
 )
 
 add_builtin(
@@ -1832,9 +2226,9 @@ add_builtin(
     value_type=uint32,
     group="Random",
     doc="""Generate a random sample from a Poisson distribution.
-    
-    :param state: RNG state
-    :param lam: The expected value of the distribution""",
+
+   :param state: RNG state
+   :param lam: The expected value of the distribution""",
 )
 
 add_builtin(
@@ -1842,28 +2236,28 @@ add_builtin(
     input_types={"state": uint32, "x": float},
     value_type=float,
     group="Random",
-    doc="Non-periodic Perlin-style noise in 1d.",
+    doc="Non-periodic Perlin-style noise in 1D.",
 )
 add_builtin(
     "noise",
     input_types={"state": uint32, "xy": vec2},
     value_type=float,
     group="Random",
-    doc="Non-periodic Perlin-style noise in 2d.",
+    doc="Non-periodic Perlin-style noise in 2D.",
 )
 add_builtin(
     "noise",
     input_types={"state": uint32, "xyz": vec3},
     value_type=float,
     group="Random",
-    doc="Non-periodic Perlin-style noise in 3d.",
+    doc="Non-periodic Perlin-style noise in 3D.",
 )
 add_builtin(
     "noise",
     input_types={"state": uint32, "xyzt": vec4},
     value_type=float,
     group="Random",
-    doc="Non-periodic Perlin-style noise in 4d.",
+    doc="Non-periodic Perlin-style noise in 4D.",
 )
 
 add_builtin(
@@ -1871,33 +2265,34 @@ add_builtin(
     input_types={"state": uint32, "x": float, "px": int},
     value_type=float,
     group="Random",
-    doc="Periodic Perlin-style noise in 1d.",
+    doc="Periodic Perlin-style noise in 1D.",
 )
 add_builtin(
     "pnoise",
     input_types={"state": uint32, "xy": vec2, "px": int, "py": int},
     value_type=float,
     group="Random",
-    doc="Periodic Perlin-style noise in 2d.",
+    doc="Periodic Perlin-style noise in 2D.",
 )
 add_builtin(
     "pnoise",
     input_types={"state": uint32, "xyz": vec3, "px": int, "py": int, "pz": int},
     value_type=float,
     group="Random",
-    doc="Periodic Perlin-style noise in 3d.",
+    doc="Periodic Perlin-style noise in 3D.",
 )
 add_builtin(
     "pnoise",
     input_types={"state": uint32, "xyzt": vec4, "px": int, "py": int, "pz": int, "pt": int},
     value_type=float,
     group="Random",
-    doc="Periodic Perlin-style noise in 4d.",
+    doc="Periodic Perlin-style noise in 4D.",
 )
 
 add_builtin(
     "curlnoise",
-    input_types={"state": uint32, "xy": vec2},
+    input_types={"state": uint32, "xy": vec2, "octaves": uint32, "lacunarity": float, "gain": float},
+    defaults={"octaves": 1, "lacunarity": 2.0, "gain": 0.5},
     value_type=vec2,
     group="Random",
     doc="Divergence-free vector field based on the gradient of a Perlin noise function.",
@@ -1905,7 +2300,8 @@ add_builtin(
 )
 add_builtin(
     "curlnoise",
-    input_types={"state": uint32, "xyz": vec3},
+    input_types={"state": uint32, "xyz": vec3, "octaves": uint32, "lacunarity": float, "gain": float},
+    defaults={"octaves": 1, "lacunarity": 2.0, "gain": 0.5},
     value_type=vec3,
     group="Random",
     doc="Divergence-free vector field based on the curl of three Perlin noise functions.",
@@ -1913,7 +2309,8 @@ add_builtin(
 )
 add_builtin(
     "curlnoise",
-    input_types={"state": uint32, "xyzt": vec4},
+    input_types={"state": uint32, "xyzt": vec4, "octaves": uint32, "lacunarity": float, "gain": float},
+    defaults={"octaves": 1, "lacunarity": 2.0, "gain": 0.5},
     value_type=vec3,
     group="Random",
     doc="Divergence-free vector field based on the curl of three Perlin noise functions.",
@@ -1927,7 +2324,7 @@ add_builtin(
     namespace="",
     variadic=True,
     group="Utility",
-    doc="Allows printing formatted strings, using C-style format specifiers.",
+    doc="Allows printing formatted strings using C-style format specifiers.",
 )
 
 add_builtin("print", input_types={"value": Any}, doc="Print variable to stdout", export=False, group="Utility")
@@ -1947,9 +2344,12 @@ add_builtin(
     "tid",
     input_types={},
     value_type=int,
+    export=False,
     group="Utility",
-    doc="""Return the current thread index. Note that this is the *global* index of the thread in the range [0, dim) 
-   where dim is the parameter passed to kernel launch.""",
+    doc="""Return the current thread index for a 1D kernel launch. Note that this is the *global* index of the thread in the range [0, dim)
+   where dim is the parameter passed to kernel launch. This function may not be called from user-defined Warp functions.""",
+    namespace="",
+    native_func="builtin_tid1d",
 )
 
 add_builtin(
@@ -1957,7 +2357,10 @@ add_builtin(
     input_types={},
     value_type=[int, int],
     group="Utility",
-    doc="""Return the current thread indices for a 2d kernel launch. Use ``i,j = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.""",
+    doc="""Return the current thread indices for a 2D kernel launch. Use ``i,j = wp.tid()`` syntax to retrieve the
+   coordinates inside the kernel thread grid. This function may not be called from user-defined Warp functions.""",
+    namespace="",
+    native_func="builtin_tid2d",
 )
 
 add_builtin(
@@ -1965,7 +2368,10 @@ add_builtin(
     input_types={},
     value_type=[int, int, int],
     group="Utility",
-    doc="""Return the current thread indices for a 3d kernel launch. Use ``i,j,k = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.""",
+    doc="""Return the current thread indices for a 3D kernel launch. Use ``i,j,k = wp.tid()`` syntax to retrieve the
+   coordinates inside the kernel thread grid. This function may not be called from user-defined Warp functions.""",
+    namespace="",
+    native_func="builtin_tid3d",
 )
 
 add_builtin(
@@ -1973,42 +2379,60 @@ add_builtin(
     input_types={},
     value_type=[int, int, int, int],
     group="Utility",
-    doc="""Return the current thread indices for a 4d kernel launch. Use ``i,j,k,l = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.""",
+    doc="""Return the current thread indices for a 4D kernel launch. Use ``i,j,k,l = wp.tid()`` syntax to retrieve the
+   coordinates inside the kernel thread grid. This function may not be called from user-defined Warp functions.""",
+    namespace="",
+    native_func="builtin_tid4d",
 )
 
 
-add_builtin("copy", variadic=True, hidden=True, export=False, group="Utility")
+add_builtin(
+    "copy",
+    input_types={"value": Any},
+    value_func=lambda arg_types, kwds, _: arg_types[0],
+    hidden=True,
+    export=False,
+    group="Utility",
+)
+add_builtin("assign", variadic=True, hidden=True, export=False, group="Utility")
 add_builtin(
     "select",
     input_types={"cond": bool, "arg1": Any, "arg2": Any},
+    value_func=lambda arg_types, kwds, _: arg_types[1],
+    doc="Select between two arguments, if ``cond`` is ``False`` then return ``arg1``, otherwise return ``arg2``",
+    group="Utility",
+)
+add_builtin(
+    "select",
+    input_types={"cond": builtins.bool, "arg1": Any, "arg2": Any},
     value_func=lambda args, kwds, _: args[1].type,
-    doc="Select between two arguments, if cond is false then return ``arg1``, otherwise return ``arg2``",
+    doc="Select between two arguments, if ``cond`` is ``False`` then return ``arg1``, otherwise return ``arg2``",
     group="Utility",
 )
 for t in int_types:
     add_builtin(
         "select",
         input_types={"cond": t, "arg1": Any, "arg2": Any},
-        value_func=lambda args, kwds, _: args[1].type,
-        doc="Select between two arguments, if cond is false then return ``arg1``, otherwise return ``arg2``",
+        value_func=lambda arg_types, kwds, _: arg_types[1],
+        doc="Select between two arguments, if ``cond`` is ``False`` then return ``arg1``, otherwise return ``arg2``",
         group="Utility",
     )
 add_builtin(
     "select",
     input_types={"arr": array(dtype=Any), "arg1": Any, "arg2": Any},
-    value_func=lambda args, kwds, _: args[1].type,
-    doc="Select between two arguments, if array is null then return ``arg1``, otherwise return ``arg2``",
+    value_func=lambda arg_types, kwds, _: arg_types[1],
+    doc="Select between two arguments, if ``arr`` is null then return ``arg1``, otherwise return ``arg2``",
     group="Utility",
 )
 
 
-# does argument checking and type propagation for load()
-def load_value_func(args, kwds, _):
-    if not is_array(args[0].type):
+# does argument checking and type propagation for address()
+def address_value_func(arg_types, kwds, _):
+    if not is_array(arg_types[0]):
         raise RuntimeError("load() argument 0 must be an array")
 
-    num_indices = len(args[1:])
-    num_dims = args[0].type.ndim
+    num_indices = len(arg_types[1:])
+    num_dims = arg_types[0].ndim
 
     if num_indices < num_dims:
         raise RuntimeError(
@@ -2021,21 +2445,21 @@ def load_value_func(args, kwds, _):
         )
 
     # check index types
-    for a in args[1:]:
-        if type_is_int(a.type) == False:
-            raise RuntimeError(f"load() index arguments must be of integer type, got index of type {a.type}")
+    for t in arg_types[1:]:
+        if not type_is_int(t):
+            raise RuntimeError(f"address() index arguments must be of integer type, got index of type {t}")
 
-    return args[0].type.dtype
+    return Reference(arg_types[0].dtype)
 
 
 # does argument checking and type propagation for view()
-def view_value_func(args, kwds, _):
-    if not is_array(args[0].type):
+def view_value_func(arg_types, kwds, _):
+    if not is_array(arg_types[0]):
         raise RuntimeError("view() argument 0 must be an array")
 
     # check array dim big enough to support view
-    num_indices = len(args[1:])
-    num_dims = args[0].type.ndim
+    num_indices = len(arg_types[1:])
+    num_dims = arg_types[0].ndim
 
     if num_indices >= num_dims:
         raise RuntimeError(
@@ -2043,27 +2467,28 @@ def view_value_func(args, kwds, _):
         )
 
     # check index types
-    for a in args[1:]:
-        if type_is_int(a.type) == False:
-            raise RuntimeError(f"view() index arguments must be of integer type, got index of type {a.type}")
+    for t in arg_types[1:]:
+        if not type_is_int(t):
+            raise RuntimeError(f"view() index arguments must be of integer type, got index of type {t}")
 
     # create an array view with leading dimensions removed
-    import copy
-
-    view_type = copy.copy(args[0].type)
-    view_type.ndim -= num_indices
-
-    return view_type
+    dtype = arg_types[0].dtype
+    ndim = num_dims - num_indices
+    if isinstance(arg_types[0], (fabricarray, indexedfabricarray)):
+        # fabric array of arrays: return array attribute as a regular array
+        return array(dtype=dtype, ndim=ndim)
+    else:
+        return type(arg_types[0])(dtype=dtype, ndim=ndim)
 
 
-# does argument checking and type propagation for store()
-def store_value_func(args, kwds, _):
+# does argument checking and type propagation for array_store()
+def array_store_value_func(arg_types, kwds, _):
     # check target type
-    if not is_array(args[0].type):
-        raise RuntimeError("store() argument 0 must be an array")
+    if not is_array(arg_types[0]):
+        raise RuntimeError("array_store() argument 0 must be an array")
 
-    num_indices = len(args[1:-1])
-    num_dims = args[0].type.ndim
+    num_indices = len(arg_types[1:-1])
+    num_dims = arg_types[0].ndim
 
     # if this happens we should have generated a view instead of a load during code gen
     if num_indices < num_dims:
@@ -2075,31 +2500,63 @@ def store_value_func(args, kwds, _):
         )
 
     # check index types
-    for a in args[1:-1]:
-        if type_is_int(a.type) == False:
-            raise RuntimeError(f"store() index arguments must be of integer type, got index of type {a.type}")
+    for t in arg_types[1:-1]:
+        if not type_is_int(t):
+            raise RuntimeError(f"array_store() index arguments must be of integer type, got index of type {t}")
 
     # check value type
-    if not types_equal(args[-1].type, args[0].type.dtype):
+    if not types_equal(arg_types[-1], arg_types[0].dtype):
         raise RuntimeError(
-            f"store() value argument type ({args[2].type}) must be of the same type as the array ({args[0].type.dtype})"
+            f"array_store() value argument type ({arg_types[2]}) must be of the same type as the array ({arg_types[0].dtype})"
         )
 
     return None
 
 
-add_builtin("load", variadic=True, hidden=True, value_func=load_value_func, group="Utility")
+# does argument checking for store()
+def store_value_func(arg_types, kwds, _):
+    # we already stripped the Reference from the argument type prior to this call
+    if not types_equal(arg_types[0], arg_types[1]):
+        raise RuntimeError(f"store() value argument type ({arg_types[1]}) must be of the same type as the reference")
+
+    return None
+
+
+# does type propagation for load()
+def load_value_func(arg_types, kwds, _):
+    # we already stripped the Reference from the argument type prior to this call
+    return arg_types[0]
+
+
+add_builtin("address", variadic=True, hidden=True, value_func=address_value_func, group="Utility")
 add_builtin("view", variadic=True, hidden=True, value_func=view_value_func, group="Utility")
-add_builtin("store", variadic=True, hidden=True, value_func=store_value_func, skip_replay=True, group="Utility")
+add_builtin(
+    "array_store", variadic=True, hidden=True, value_func=array_store_value_func, skip_replay=True, group="Utility"
+)
+add_builtin(
+    "store",
+    input_types={"address": Reference, "value": Any},
+    hidden=True,
+    value_func=store_value_func,
+    skip_replay=True,
+    group="Utility",
+)
+add_builtin(
+    "load",
+    input_types={"address": Reference},
+    hidden=True,
+    value_func=load_value_func,
+    group="Utility",
+)
 
 
-def atomic_op_value_func(args, kwds, _):
+def atomic_op_value_func(arg_types, kwds, _):
     # check target type
-    if not is_array(args[0].type):
+    if not is_array(arg_types[0]):
         raise RuntimeError("atomic() operation argument 0 must be an array")
 
-    num_indices = len(args[1:-1])
-    num_dims = args[0].type.ndim
+    num_indices = len(arg_types[1:-1])
+    num_dims = arg_types[0].ndim
 
     # if this happens we should have generated a view instead of a load during code gen
     if num_indices < num_dims:
@@ -2111,18 +2568,16 @@ def atomic_op_value_func(args, kwds, _):
         )
 
     # check index types
-    for a in args[1:-1]:
-        if type_is_int(a.type) == False:
-            raise RuntimeError(
-                f"atomic() operation index arguments must be of integer type, got index of type {a.type}"
-            )
+    for t in arg_types[1:-1]:
+        if not type_is_int(t):
+            raise RuntimeError(f"atomic() operation index arguments must be of integer type, got index of type {t}")
 
-    if not types_equal(args[-1].type, args[0].type.dtype):
+    if not types_equal(arg_types[-1], arg_types[0].dtype):
         raise RuntimeError(
-            f"atomic() value argument ({args[-1].type}) must be of the same type as the array ({args[0].type.dtype})"
+            f"atomic() value argument ({arg_types[-1]}) must be of the same type as the array ({arg_types[0].dtype})"
         )
 
-    return args[0].type.dtype
+    return arg_types[0].dtype
 
 
 for array_type in array_types:
@@ -2134,7 +2589,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto the array at location given by index.",
+        doc="Atomically add ``value`` onto ``a[i]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2143,7 +2598,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto the array at location given by indices.",
+        doc="Atomically add ``value`` onto ``a[i,j]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2152,7 +2607,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto the array at location given by indices.",
+        doc="Atomically add ``value`` onto ``a[i,j,k]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2161,7 +2616,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically add ``value`` onto the array at location given by indices.",
+        doc="Atomically add ``value`` onto ``a[i,j,k,l]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2171,7 +2626,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto the array at location given by index.",
+        doc="Atomically subtract ``value`` onto ``a[i]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2180,7 +2635,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto the array at location given by indices.",
+        doc="Atomically subtract ``value`` onto ``a[i,j]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2189,7 +2644,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto the array at location given by indices.",
+        doc="Atomically subtract ``value`` onto ``a[i,j,k]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2198,7 +2653,7 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Atomically subtract ``value`` onto the array at location given by indices.",
+        doc="Atomically subtract ``value`` onto ``a[i,j,k,l]``.",
         group="Utility",
         skip_replay=True,
     )
@@ -2208,7 +2663,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the minimum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the minimum of ``value`` and ``a[i]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2217,7 +2673,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the minimum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the minimum of ``value`` and ``a[i,j]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2226,7 +2683,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the minimum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the minimum of ``value`` and ``a[i,j,k]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2235,7 +2693,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the minimum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the minimum of ``value`` and ``a[i,j,k,l]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2245,7 +2704,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the maximum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the maximum of ``value`` and ``a[i]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2254,7 +2714,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the maximum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the maximum of ``value`` and ``a[i,j]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2263,7 +2724,8 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the maximum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the maximum of ``value`` and ``a[i,j,k]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
@@ -2272,26 +2734,27 @@ for array_type in array_types:
         hidden=hidden,
         input_types={"a": array_type(dtype=Any), "i": int, "j": int, "k": int, "l": int, "value": Any},
         value_func=atomic_op_value_func,
-        doc="Compute the maximum of ``value`` and ``array[index]`` and atomically update the array. Note that for vectors and matrices the operation is only atomic on a per-component basis.",
+        doc="Compute the maximum of ``value`` and ``a[i,j,k,l]`` and atomically update the array.\n\n"
+        "Note that for vectors and matrices the operation is only atomic on a per-component basis.",
         group="Utility",
         skip_replay=True,
     )
 
 
 # used to index into builtin types, i.e.: y = vec3[1]
-def index_value_func(args, kwds, _):
-    return args[0].type._wp_scalar_type_
+def index_value_func(arg_types, kwds, _):
+    return arg_types[0]._wp_scalar_type_
 
 
 add_builtin(
-    "index",
+    "extract",
     input_types={"a": vector(length=Any, dtype=Scalar), "i": int},
     value_func=index_value_func,
     hidden=True,
     group="Utility",
 )
 add_builtin(
-    "index",
+    "extract",
     input_types={"a": quaternion(dtype=Scalar), "i": int},
     value_func=index_value_func,
     hidden=True,
@@ -2299,14 +2762,14 @@ add_builtin(
 )
 
 add_builtin(
-    "index",
+    "extract",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int},
-    value_func=lambda args, kwds, _: vector(length=args[0].type._shape_[1], dtype=args[0].type._wp_scalar_type_),
+    value_func=lambda arg_types, kwds, _: vector(length=arg_types[0]._shape_[1], dtype=arg_types[0]._wp_scalar_type_),
     hidden=True,
     group="Utility",
 )
 add_builtin(
-    "index",
+    "extract",
     input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int},
     value_func=index_value_func,
     hidden=True,
@@ -2314,77 +2777,66 @@ add_builtin(
 )
 
 add_builtin(
-    "index",
+    "extract",
     input_types={"a": transformation(dtype=Scalar), "i": int},
     value_func=index_value_func,
     hidden=True,
     group="Utility",
 )
 
-add_builtin("index", input_types={"s": shape_t, "i": int}, value_type=int, hidden=True, group="Utility")
+add_builtin("extract", input_types={"s": shape_t, "i": int}, value_type=int, hidden=True, group="Utility")
 
 
-def vector_indexset_element_value_func(args, kwds, _):
-    vec = args[0]
-    index = args[1]
-    value = args[2]
+def vector_indexref_element_value_func(arg_types, kwds, _):
+    vec_type = arg_types[0]
+    # index_type = arg_types[1]
+    value_type = vec_type._wp_scalar_type_
 
-    if value.type is not vec.type._wp_scalar_type_:
-        raise RuntimeError(
-            f"Trying to assign type '{type_repr(value.type)}' to element of a vector with type '{type_repr(vec.type)}'"
-        )
-
-    return None
+    return Reference(value_type)
 
 
-# implements vector[index] = value
+# implements &vector[index]
 add_builtin(
-    "indexset",
-    input_types={"a": vector(length=Any, dtype=Scalar), "i": int, "value": Scalar},
-    value_func=vector_indexset_element_value_func,
+    "index",
+    input_types={"a": vector(length=Any, dtype=Scalar), "i": int},
+    value_func=vector_indexref_element_value_func,
+    hidden=True,
+    group="Utility",
+    skip_replay=True,
+)
+# implements &(*vector)[index]
+add_builtin(
+    "indexref",
+    input_types={"a": Reference, "i": int},
+    value_func=vector_indexref_element_value_func,
     hidden=True,
     group="Utility",
     skip_replay=True,
 )
 
 
-def matrix_indexset_element_value_func(args, kwds, _):
-    mat = args[0]
-    row = args[1]
-    col = args[2]
-    value = args[3]
-
-    if value.type is not mat.type._wp_scalar_type_:
-        raise RuntimeError(
-            f"Trying to assign type '{type_repr(value.type)}' to element of a matrix with type '{type_repr(mat.type)}'"
-        )
-
-    return None
-
+def matrix_indexref_element_value_func(arg_types, kwds, _):
+    mat_type = arg_types[0]
+    # row_type = arg_types[1]
+    # col_type = arg_types[2]
+    value_type = mat_type._wp_scalar_type_
 
-def matrix_indexset_row_value_func(args, kwds, _):
-    mat = args[0]
-    row = args[1]
-    value = args[2]
+    return Reference(value_type)
 
-    if value.type._shape_[0] != mat.type._shape_[1]:
-        raise RuntimeError(
-            f"Trying to assign vector with length {value.type._length} to matrix with shape {mat.type._shape}, vector length must match the number of matrix columns."
-        )
 
-    if value.type._wp_scalar_type_ is not mat.type._wp_scalar_type_:
-        raise RuntimeError(
-            f"Trying to assign vector of type '{type_repr(value.type)}' to row of matrix of type '{type_repr(mat.type)}'"
-        )
+def matrix_indexref_row_value_func(arg_types, kwds, _):
+    mat_type = arg_types[0]
+    row_type = mat_type._wp_row_type_
+    # value_type = arg_types[2]
 
-    return None
+    return Reference(row_type)
 
 
 # implements matrix[i] = row
 add_builtin(
-    "indexset",
-    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "value": vector(length=Any, dtype=Scalar)},
-    value_func=matrix_indexset_row_value_func,
+    "index",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int},
+    value_func=matrix_indexref_row_value_func,
     hidden=True,
     group="Utility",
     skip_replay=True,
@@ -2392,29 +2844,29 @@ add_builtin(
 
 # implements matrix[i,j] = scalar
 add_builtin(
-    "indexset",
-    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int, "value": Scalar},
-    value_func=matrix_indexset_element_value_func,
+    "index",
+    input_types={"a": matrix(shape=(Any, Any), dtype=Scalar), "i": int, "j": int},
+    value_func=matrix_indexref_element_value_func,
     hidden=True,
     group="Utility",
     skip_replay=True,
 )
 
-for t in scalar_types + vector_types:
+for t in scalar_types + vector_types + [builtins.bool]:
     if "vec" in t.__name__ or "mat" in t.__name__:
         continue
     add_builtin(
         "expect_eq",
         input_types={"arg1": t, "arg2": t},
         value_type=None,
-        doc="Prints an error to stdout if arg1 and arg2 are not equal",
+        doc="Prints an error to stdout if ``arg1`` and ``arg2`` are not equal",
         group="Utility",
         hidden=True,
     )
 
 
-def expect_eq_val_func(args, kwds, _):
-    if not types_equal(args[0].type, args[1].type):
+def expect_eq_val_func(arg_types, kwds, _):
+    if not types_equal(arg_types[0], arg_types[1]):
         raise RuntimeError("Can't test equality for objects with different types")
     return None
 
@@ -2423,7 +2875,7 @@ add_builtin(
     "expect_eq",
     input_types={"arg1": vector(length=Any, dtype=Scalar), "arg2": vector(length=Any, dtype=Scalar)},
     value_func=expect_eq_val_func,
-    doc="Prints an error to stdout if arg1 and arg2 are not equal",
+    doc="Prints an error to stdout if ``arg1`` and ``arg2`` are not equal",
     group="Utility",
     hidden=True,
 )
@@ -2431,7 +2883,7 @@ add_builtin(
     "expect_neq",
     input_types={"arg1": vector(length=Any, dtype=Scalar), "arg2": vector(length=Any, dtype=Scalar)},
     value_func=expect_eq_val_func,
-    doc="Prints an error to stdout if arg1 and arg2 are equal",
+    doc="Prints an error to stdout if ``arg1`` and ``arg2`` are equal",
     group="Utility",
     hidden=True,
 )
@@ -2440,7 +2892,7 @@ add_builtin(
     "expect_eq",
     input_types={"arg1": matrix(shape=(Any, Any), dtype=Scalar), "arg2": matrix(shape=(Any, Any), dtype=Scalar)},
     value_func=expect_eq_val_func,
-    doc="Prints an error to stdout if arg1 and arg2 are not equal",
+    doc="Prints an error to stdout if ``arg1`` and ``arg2`` are not equal",
     group="Utility",
     hidden=True,
 )
@@ -2448,7 +2900,7 @@ add_builtin(
     "expect_neq",
     input_types={"arg1": matrix(shape=(Any, Any), dtype=Scalar), "arg2": matrix(shape=(Any, Any), dtype=Scalar)},
     value_func=expect_eq_val_func,
-    doc="Prints an error to stdout if arg1 and arg2 are equal",
+    doc="Prints an error to stdout if ``arg1`` and ``arg2`` are equal",
     group="Utility",
     hidden=True,
 )
@@ -2457,29 +2909,30 @@ add_builtin(
     "lerp",
     input_types={"a": Float, "b": Float, "t": Float},
     value_func=sametype_value_func(Float),
-    doc="Linearly interpolate two values a and b using factor t, computed as ``a*(1-t) + b*t``",
+    doc="Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``",
     group="Utility",
 )
 add_builtin(
     "smoothstep",
     input_types={"edge0": Float, "edge1": Float, "x": Float},
     value_func=sametype_value_func(Float),
-    doc="Smoothly interpolate between two values edge0 and edge1 using a factor x, and return a result between 0 and 1 using a cubic Hermite interpolation after clamping",
+    doc="""Smoothly interpolate between two values ``edge0`` and ``edge1`` using a factor ``x``,
+   and return a result between 0 and 1 using a cubic Hermite interpolation after clamping.""",
     group="Utility",
 )
 
 
 def lerp_value_func(default):
-    def fn(args, kwds, _):
-        if args is None:
+    def fn(arg_types, kwds, _):
+        if arg_types is None:
             return default
-        scalar_type = args[-1].type
-        if not types_equal(args[0].type, args[1].type):
+        scalar_type = arg_types[-1]
+        if not types_equal(arg_types[0], arg_types[1]):
             raise RuntimeError("Can't lerp between objects with different types")
-        if args[0].type._wp_scalar_type_ != scalar_type:
+        if arg_types[0]._wp_scalar_type_ != scalar_type:
             raise RuntimeError("'t' parameter must have the same scalar type as objects you're lerping between")
 
-        return args[0].type
+        return arg_types[0]
 
     return fn
 
@@ -2488,28 +2941,28 @@ add_builtin(
     "lerp",
     input_types={"a": vector(length=Any, dtype=Float), "b": vector(length=Any, dtype=Float), "t": Float},
     value_func=lerp_value_func(vector(length=Any, dtype=Float)),
-    doc="Linearly interpolate two values a and b using factor t, computed as ``a*(1-t) + b*t``",
+    doc="Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``",
     group="Utility",
 )
 add_builtin(
     "lerp",
     input_types={"a": matrix(shape=(Any, Any), dtype=Float), "b": matrix(shape=(Any, Any), dtype=Float), "t": Float},
     value_func=lerp_value_func(matrix(shape=(Any, Any), dtype=Float)),
-    doc="Linearly interpolate two values a and b using factor t, computed as ``a*(1-t) + b*t``",
+    doc="Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``",
     group="Utility",
 )
 add_builtin(
     "lerp",
     input_types={"a": quaternion(dtype=Float), "b": quaternion(dtype=Float), "t": Float},
     value_func=lerp_value_func(quaternion(dtype=Float)),
-    doc="Linearly interpolate two values a and b using factor t, computed as ``a*(1-t) + b*t``",
+    doc="Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``",
     group="Utility",
 )
 add_builtin(
     "lerp",
     input_types={"a": transformation(dtype=Float), "b": transformation(dtype=Float), "t": Float},
     value_func=lerp_value_func(transformation(dtype=Float)),
-    doc="Linearly interpolate two values a and b using factor t, computed as ``a*(1-t) + b*t``",
+    doc="Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``",
     group="Utility",
 )
 
@@ -2519,14 +2972,14 @@ add_builtin(
     input_types={"arg1": Float, "arg2": Float, "tolerance": Float},
     defaults={"tolerance": 1.0e-6},
     value_type=None,
-    doc="Prints an error to stdout if arg1 and arg2 are not closer than tolerance in magnitude",
+    doc="Prints an error to stdout if ``arg1`` and ``arg2`` are not closer than tolerance in magnitude",
     group="Utility",
 )
 add_builtin(
     "expect_near",
     input_types={"arg1": vec3, "arg2": vec3, "tolerance": float},
     value_type=None,
-    doc="Prints an error to stdout if any element of arg1 and arg2 are not closer than tolerance in magnitude",
+    doc="Prints an error to stdout if any element of ``arg1`` and ``arg2`` are not closer than tolerance in magnitude",
     group="Utility",
 )
 
@@ -2537,7 +2990,14 @@ add_builtin(
     "lower_bound",
     input_types={"arr": array(dtype=Scalar), "value": Scalar},
     value_type=int,
-    doc="Search a sorted array for the closest element greater than or equal to value.",
+    doc="Search a sorted array ``arr`` for the closest element greater than or equal to ``value``.",
+)
+
+add_builtin(
+    "lower_bound",
+    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``.",
 )
 
 # ---------------------------------
@@ -2617,11 +3077,11 @@ add_builtin("invert", input_types={"x": Int}, value_func=sametype_value_func(Int
 
 
 def scalar_mul_value_func(default):
-    def fn(args, kwds, _):
-        if args is None:
+    def fn(arg_types, kwds, _):
+        if arg_types is None:
             return default
-        scalar = [a.type for a in args if a.type in scalar_types][0]
-        compound = [a.type for a in args if a.type not in scalar_types][0]
+        scalar = [t for t in arg_types if t in scalar_types][0]
+        compound = [t for t in arg_types if t not in scalar_types][0]
         if scalar != compound._wp_scalar_type_:
             raise RuntimeError("Object and coefficient must have the same scalar type when multiplying by scalar")
         return compound
@@ -2629,36 +3089,53 @@ def scalar_mul_value_func(default):
     return fn
 
 
-def mul_matvec_value_func(args, kwds, _):
-    if args is None:
+def mul_matvec_value_func(arg_types, kwds, _):
+    if arg_types is None:
+        return vector(length=Any, dtype=Scalar)
+
+    if arg_types[0]._wp_scalar_type_ != arg_types[1]._wp_scalar_type_:
+        raise RuntimeError(
+            f"Can't multiply matrix and vector with different types {arg_types[0]._wp_scalar_type_}, {arg_types[1]._wp_scalar_type_}"
+        )
+
+    if arg_types[0]._shape_[1] != arg_types[1]._length_:
+        raise RuntimeError(
+            f"Can't multiply matrix of shape {arg_types[0]._shape_} and vector with length {arg_types[1]._length_}"
+        )
+
+    return vector(length=arg_types[0]._shape_[0], dtype=arg_types[0]._wp_scalar_type_)
+
+
+def mul_vecmat_value_func(arg_types, kwds, _):
+    if arg_types is None:
         return vector(length=Any, dtype=Scalar)
 
-    if args[0].type._wp_scalar_type_ != args[1].type._wp_scalar_type_:
+    if arg_types[1]._wp_scalar_type_ != arg_types[0]._wp_scalar_type_:
         raise RuntimeError(
-            f"Can't multiply matrix and vector with different types {args[0].type._wp_scalar_type_}, {args[1].type._wp_scalar_type_}"
+            f"Can't multiply vector and matrix with different types {arg_types[1]._wp_scalar_type_}, {arg_types[0]._wp_scalar_type_}"
         )
 
-    if args[0].type._shape_[1] != args[1].type._length_:
+    if arg_types[1]._shape_[0] != arg_types[0]._length_:
         raise RuntimeError(
-            f"Can't multiply matrix of shape {args[0].type._shape_} and vector with length {args[1].type._length_}"
+            f"Can't multiply vector with length {arg_types[0]._length_} and matrix of shape {arg_types[1]._shape_}"
         )
 
-    return vector(length=args[0].type._shape_[0], dtype=args[0].type._wp_scalar_type_)
+    return vector(length=arg_types[1]._shape_[1], dtype=arg_types[1]._wp_scalar_type_)
 
 
-def mul_matmat_value_func(args, kwds, _):
-    if args is None:
+def mul_matmat_value_func(arg_types, kwds, _):
+    if arg_types is None:
         return matrix(length=Any, dtype=Scalar)
 
-    if args[0].type._wp_scalar_type_ != args[1].type._wp_scalar_type_:
+    if arg_types[0]._wp_scalar_type_ != arg_types[1]._wp_scalar_type_:
         raise RuntimeError(
-            f"Can't multiply matrices with different types {args[0].type._wp_scalar_type_}, {args[1].type._wp_scalar_type_}"
+            f"Can't multiply matrices with different types {arg_types[0]._wp_scalar_type_}, {arg_types[1]._wp_scalar_type_}"
         )
 
-    if args[0].type._shape_[1] != args[1].type._shape_[0]:
-        raise RuntimeError(f"Can't multiply matrix of shapes {args[0].type._shape_} and {args[1].type._shape_}")
+    if arg_types[0]._shape_[1] != arg_types[1]._shape_[0]:
+        raise RuntimeError(f"Can't multiply matrix of shapes {arg_types[0]._shape_} and {arg_types[1]._shape_}")
 
-    return matrix(shape=(args[0].type._shape_[0], args[1].type._shape_[1]), dtype=args[0].type._wp_scalar_type_)
+    return matrix(shape=(arg_types[0]._shape_[0], arg_types[1]._shape_[1]), dtype=arg_types[0]._wp_scalar_type_)
 
 
 add_builtin(
@@ -2720,6 +3197,13 @@ add_builtin(
     doc="",
     group="Operators",
 )
+add_builtin(
+    "mul",
+    input_types={"x": vector(length=Any, dtype=Scalar), "y": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_func=mul_vecmat_value_func,
+    doc="",
+    group="Operators",
+)
 add_builtin(
     "mul",
     input_types={"x": matrix(shape=(Any, Any), dtype=Scalar), "y": matrix(shape=(Any, Any), dtype=Scalar)},
@@ -2755,7 +3239,12 @@ add_builtin(
 )
 
 add_builtin(
-    "div", input_types={"x": Scalar, "y": Scalar}, value_func=sametype_value_func(Scalar), doc="", group="Operators"
+    "div",
+    input_types={"x": Scalar, "y": Scalar},
+    value_func=sametype_value_func(Scalar),
+    doc="",
+    group="Operators",
+    require_original_output_arg=True,
 )
 add_builtin(
     "div",
@@ -2764,6 +3253,13 @@ add_builtin(
     doc="",
     group="Operators",
 )
+add_builtin(
+    "div",
+    input_types={"x": Scalar, "y": vector(length=Any, dtype=Scalar)},
+    value_func=scalar_mul_value_func(vector(length=Any, dtype=Scalar)),
+    doc="",
+    group="Operators",
+)
 add_builtin(
     "div",
     input_types={"x": matrix(shape=(Any, Any), dtype=Scalar), "y": Scalar},
@@ -2771,6 +3267,13 @@ add_builtin(
     doc="",
     group="Operators",
 )
+add_builtin(
+    "div",
+    input_types={"x": Scalar, "y": matrix(shape=(Any, Any), dtype=Scalar)},
+    value_func=scalar_mul_value_func(matrix(shape=(Any, Any), dtype=Scalar)),
+    doc="",
+    group="Operators",
+)
 add_builtin(
     "div",
     input_types={"x": quaternion(dtype=Scalar), "y": Scalar},
@@ -2778,6 +3281,13 @@ add_builtin(
     doc="",
     group="Operators",
 )
+add_builtin(
+    "div",
+    input_types={"x": Scalar, "y": quaternion(dtype=Scalar)},
+    value_func=scalar_mul_value_func(quaternion(dtype=Scalar)),
+    doc="",
+    group="Operators",
+)
 
 add_builtin(
     "floordiv",
@@ -2832,9 +3342,9 @@ add_builtin(
     group="Operators",
 )
 
-add_builtin("unot", input_types={"b": bool}, value_type=bool, doc="", group="Operators")
+add_builtin("unot", input_types={"b": builtins.bool}, value_type=builtins.bool, doc="", group="Operators")
 for t in int_types:
-    add_builtin("unot", input_types={"b": t}, value_type=bool, doc="", group="Operators")
+    add_builtin("unot", input_types={"b": t}, value_type=builtins.bool, doc="", group="Operators")
 
 
-add_builtin("unot", input_types={"a": array(dtype=Any)}, value_type=bool, doc="", group="Operators")
+add_builtin("unot", input_types={"a": array(dtype=Any)}, value_type=builtins.bool, doc="", group="Operators")