warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.1__py3-none-macosx_10_13_universal2.whl

warp/types.py

@@ -20,7 +20,21 @@ import ctypes
 import inspect
 import struct
 import zlib
-from typing import Any, Callable, Generic, List, Literal, NamedTuple, Optional, Sequence, Tuple, TypeVar, Union
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    List,
+    Literal,
+    NamedTuple,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
 
 import numpy as np
 import numpy.typing as npt
@@ -56,7 +70,9 @@ class Transformation(Generic[Float]):
 
 
 class Array(Generic[DType]):
-    pass
+    device: Optional[warp.context.Device]
+    dtype: type
+    size: int
 
 
 int_tuple_type_hints = {
@@ -262,7 +278,12 @@ def vector(length, dtype):
         def __str__(self):
             return f"[{', '.join(map(str, self))}]"
 
+        def __repr__(self):
+            return f"{type_repr(self)}([{', '.join(map(repr, self))}])"
+
         def __eq__(self, other):
+            if self._length_ != len(other):
+                return False
             for i in range(self._length_):
                 if self[i] != other[i]:
                     return False
@@ -405,7 +426,11 @@ def matrix(shape, dtype):
             return "[" + ",\n ".join(row_str) + "]"
 
         def __eq__(self, other):
+            if self._shape_[0] != len(other):
+                return False
             for i in range(self._shape_[0]):
+                if self._shape_[1] != len(other[i]):
+                    return False
                 for j in range(self._shape_[1]):
                     if self[i][j] != other[i][j]:
                         return False
@@ -1139,7 +1164,7 @@ ARRAY_TYPE_FABRIC_INDEXED = 3
 class launch_bounds_t(ctypes.Structure):
     _fields_ = [("shape", ctypes.c_int32 * LAUNCH_MAX_DIMS), ("ndim", ctypes.c_int32), ("size", ctypes.c_size_t)]
 
-    def __init__(self, shape):
+    def __init__(self, shape: Union[int, Sequence[int]]):
         if isinstance(shape, int):
             # 1d launch
             self.ndim = 1
@@ -1260,7 +1285,7 @@ _type_size_cache = {
 }
 
 
-def type_size_in_bytes(dtype):
+def type_size_in_bytes(dtype: type) -> int:
     size = _type_size_cache.get(dtype)
 
     if size is None:
@@ -1279,7 +1304,7 @@ def type_size_in_bytes(dtype):
     return size
 
 
-def type_to_warp(dtype):
+def type_to_warp(dtype: type) -> type:
     if dtype == float:
         return float32
     elif dtype == int:
@@ -1290,7 +1315,7 @@ def type_to_warp(dtype):
         return dtype
 
 
-def type_typestr(dtype):
+def type_typestr(dtype: type) -> str:
     if dtype == bool:
         return "|b1"
     elif dtype == float16:
@@ -1323,22 +1348,67 @@ def type_typestr(dtype):
         raise Exception("Unknown ctype")
 
 
+def scalar_short_name(t):
+    if t == float32:
+        return "f"
+    elif t == float64:
+        return "d"
+    elif t == int8:
+        return "b"
+    elif t == int16:
+        return "s"
+    elif t == int32:
+        return "i"
+    elif t == int64:
+        return "l"
+    elif t == uint8:
+        return "ub"
+    elif t == uint16:
+        return "us"
+    elif t == uint32:
+        return "ui"
+    elif t == uint64:
+        return "ul"
+    return None
+
+
 # converts any known type to a human readable string, good for error messages, reporting etc
 def type_repr(t):
     if is_array(t):
-        return str(f"array(ndim={t.ndim}, dtype={t.dtype})")
+        if t.device is None:
+            # array is used as a type annotation - display ndim instead of shape
+            return f"array(ndim={t.ndim}, dtype={type_repr(t.dtype)})"
+        return f"array(shape={t.shape}, dtype={type_repr(t.dtype)})"
     if is_tile(t):
-        return str(f"tile(dtype={t.dtype}, shape={t.shape}")
-    if type_is_vector(t):
-        return str(f"vector(length={t._shape_[0]}, dtype={t._wp_scalar_type_})")
-    if type_is_matrix(t):
-        return str(f"matrix(shape=({t._shape_[0]}, {t._shape_[1]}), dtype={t._wp_scalar_type_})")
+        return f"tile(shape={t.shape}, dtype={type_repr(t.dtype)})"
     if isinstance(t, warp.codegen.Struct):
         return type_repr(t.cls)
+    sn = None
+    if hasattr(t, "_wp_scalar_type_"):
+        sn = scalar_short_name(t._wp_scalar_type_)
+    if type_is_transformation(t):
+        if sn is not None:
+            return f"transform{sn}"
+        return f"transform(dtype={type_repr(t._wp_scalar_type_)})"
+    if type_is_quaternion(t):
+        if sn is not None:
+            return f"quat{sn}"
+        return f"quat(dtype={type_repr(t._wp_scalar_type_)})"
+    if type_is_vector(t):
+        if sn is not None and t._shape_[0] <= 4:
+            return f"vec{t._shape_[0]}{sn}"
+        return f"vector(length={t._shape_[0]}, dtype={type_repr(t._wp_scalar_type_)})"
+    if type_is_matrix(t):
+        if sn is not None and t._shape_[0] <= 4 and t._shape_[1] <= 4:
+            return f"mat{t._shape_[0]}{t._shape_[1]}({sn})"
+        return f"matrix(shape=({t._shape_[0]}, {t._shape_[1]}), dtype={type_repr(t._wp_scalar_type_)})"
     if t in scalar_types:
         return t.__name__
 
-    name = getattr(t, "__qualname__", t.__name__)
+    name = getattr(t, "__name__", None)
+    if name is None:
+        return repr(t)
+    name = getattr(t, "__qualname__", name)
     return t.__module__ + "." + name
 
 
@@ -1376,33 +1446,33 @@ def type_is_transformation(t):
     return getattr(t, "_wp_generic_type_hint_", None) is Transformation
 
 
-value_types = (int, float, builtins.bool) + scalar_types
+value_types = (int, float, builtins.bool) + scalar_and_bool_types
 
 
 # returns true for all value types (int, float, bool, scalars, vectors, matrices)
-def type_is_value(x):
+def type_is_value(x: Any) -> builtins.bool:
     return x in value_types or hasattr(x, "_wp_scalar_type_")
 
 
 # equivalent of the above but for values
-def is_int(x):
+def is_int(x: Any) -> builtins.bool:
     return type_is_int(type(x))
 
 
-def is_float(x):
+def is_float(x: Any) -> builtins.bool:
     return type_is_float(type(x))
 
 
-def is_value(x):
+def is_value(x: Any) -> builtins.bool:
     return type_is_value(type(x))
 
 
-# returns true if the passed *instance* is one of the array types
-def is_array(a):
+def is_array(a) -> builtins.bool:
+    """Return true if the passed *instance* is one of the array types."""
     return isinstance(a, array_types)
 
 
-def scalars_equal(a, b, match_generic):
+def scalars_equal(a, b, match_generic=False):
     # convert to canonical types
     if a == float:
         a = float32
@@ -1465,21 +1535,21 @@ def types_equal(a, b, match_generic=False):
             if a_length is None or b_length is None or a_length == b_length:
                 return True
 
-    a_origin = warp.codegen.get_type_origin(a)
-    b_origin = warp.codegen.get_type_origin(b)
+    a_origin = get_origin(a)
+    b_origin = get_origin(b)
     if a_origin is tuple and b_origin is tuple:
-        a_args = warp.codegen.get_type_args(a)
-        b_args = warp.codegen.get_type_args(b)
+        a_args = get_args(a)
+        b_args = get_args(b)
         if len(a_args) == len(b_args) and all(
             scalars_equal(x, y, match_generic=match_generic) for x, y in zip(a_args, b_args)
         ):
             return True
     elif a_origin is tuple and isinstance(b, Sequence):
-        a_args = warp.codegen.get_type_args(a)
+        a_args = get_args(a)
         if len(a_args) == len(b) and all(scalars_equal(x, y, match_generic=match_generic) for x, y in zip(a_args, b)):
             return True
     elif b_origin is tuple and isinstance(a, Sequence):
-        b_args = warp.codegen.get_type_args(b)
+        b_args = get_args(b)
         if len(b_args) == len(a) and all(scalars_equal(x, y, match_generic=match_generic) for x, y in zip(b_args, a)):
             return True
 
@@ -1600,7 +1670,7 @@ def array_ctype_from_interface(interface: dict, dtype=None, owner=None):
     return array_ctype
 
 
-class array(Array):
+class array(Array[DType]):
     """A fixed-size multi-dimensional array containing values of the same type.
 
     Attributes:
@@ -1629,21 +1699,21 @@ class array(Array):
 
     def __init__(
         self,
-        data: Optional[Union[List, Tuple, npt.NDArray]] = None,
-        dtype: Union[DType, Any] = Any,
-        shape: Optional[Tuple[int, ...]] = None,
+        data: Union[List, Tuple, npt.NDArray, None] = None,
+        dtype: Any = Any,
+        shape: Union[int, Tuple[int, ...], List[int], None] = None,
         strides: Optional[Tuple[int, ...]] = None,
         length: Optional[int] = None,
         ptr: Optional[int] = None,
         capacity: Optional[int] = None,
         device=None,
-        pinned: bool = False,
-        copy: bool = True,
-        owner: bool = False,  # deprecated - pass deleter instead
+        pinned: builtins.bool = False,
+        copy: builtins.bool = True,
+        owner: builtins.bool = False,  # deprecated - pass deleter instead
         deleter: Optional[Callable[[int, int], None]] = None,
         ndim: Optional[int] = None,
         grad: Optional[array] = None,
-        requires_grad: bool = False,
+        requires_grad: builtins.bool = False,
     ):
         """Constructs a new Warp array object
 
@@ -2219,6 +2289,9 @@ class array(Array):
         else:
             return str(self.numpy())
 
+    def __repr__(self):
+        return type_repr(self)
+
     def __getitem__(self, key):
         if isinstance(key, int):
             if self.ndim == 1:
@@ -2939,7 +3012,7 @@ def from_ipc_handle(
 
 # A base class for non-contiguous arrays, providing the implementation of common methods like
 # contiguous(), to(), numpy(), list(), assign(), zero_(), and fill_().
-class noncontiguous_array_base(Generic[T]):
+class noncontiguous_array_base(Array[T]):
     def __init__(self, array_type_id):
         self.type_id = array_type_id
         self.is_contiguous = False
@@ -3036,12 +3109,18 @@ def check_index_array(indices, expected_device):
         raise ValueError(f"Index array device ({indices.device} does not match data array device ({expected_device}))")
 
 
-class indexedarray(noncontiguous_array_base[T]):
+class indexedarray(noncontiguous_array_base):
     # member attributes available during code-gen (e.g.: d = arr.shape[0])
     # (initialized when needed)
     _vars = None
 
-    def __init__(self, data: array = None, indices: Union[array, List[array]] = None, dtype=None, ndim=None):
+    def __init__(
+        self,
+        data: Optional[array] = None,
+        indices: Union[array, List[array], None] = None,
+        dtype=None,
+        ndim: Optional[int] = None,
+    ):
         super().__init__(ARRAY_TYPE_INDEXED)
 
         # canonicalize types
@@ -3232,7 +3311,7 @@ class Tile:
                 return f"wp::tile_alloc_empty<{Var.type_to_ctype(self.dtype)},wp::tile_shape_t<{','.join(map(str, self.shape))}>,{'true' if requires_grad else 'false'}>()"
             else:
                 # tile will be initialized by another call, e.g.: tile_transpose()
-                return "NULL"
+                return "nullptr"
 
     # return total tile size in bytes
     def size_in_bytes(self):
@@ -3634,7 +3713,7 @@ class Volume:
         instance.id = None
         return instance
 
-    def __init__(self, data: array, copy: bool = True):
+    def __init__(self, data: array, copy: builtins.bool = True):
         """Class representing a sparse grid.
 
         Args:
@@ -4361,6 +4440,15 @@ class Volume:
         translation_buf = (ctypes.c_float * 3)(translation[0], translation[1], translation[2])
         return transform_buf, translation_buf
 
+    # nanovdb types for which we instantiate the grid builder
+    # Should be in sync with WP_VOLUME_BUILDER_INSTANTIATE_TYPES in volume_builder.h
+    _supported_allocation_types = [
+        "int32",
+        "float",
+        "Vec3f",
+        "Vec4f",
+    ]
+
     @classmethod
     def allocate_by_tiles(
         cls,
@@ -4388,7 +4476,8 @@ class Volume:
                 or a floating point scalar type (2D N-by-3 array of :class:`warp.float32` or 1D array of `warp.vec3f` values), indicating world space positions.
                 Repeated points per tile are allowed and will be efficiently deduplicated.
             voxel_size (float or array-like): Voxel size(s) of the new volume. Ignored if `transform` is given.
-            bg_value (array-like, float, int or None): Value of unallocated voxels of the volume, also defines the volume's type. A :class:`warp.vec3` volume is created if this is `array-like`, an index volume will be created if `bg_value` is ``None``.
+            bg_value (array-like, scalar or None): Value of unallocated voxels of the volume, also defines the volume's type. An index volume will be created if `bg_value` is ``None``.
+              Other supported grid types are `int`, `float`, `vec3f`, and `vec4f`.
             translation (array-like): Translation between the index and world spaces.
             transform (array-like): Linear transform between the index and world spaces. If ``None``, deduced from `voxel_size`.
             device (Devicelike): The CUDA device to create the volume on, e.g.: "cuda" or "cuda:0".
@@ -4420,35 +4509,47 @@ class Volume:
                 translation_buf,
                 in_world_space,
             )
-        elif hasattr(bg_value, "__len__"):
-            volume.id = volume.runtime.core.volume_v_from_tiles_device(
-                volume.device.context,
-                ctypes.c_void_p(tile_points.ptr),
-                tile_points.shape[0],
-                transform_buf,
-                translation_buf,
-                in_world_space,
-                (ctypes.c_float * 3)(bg_value[0], bg_value[1], bg_value[2]),
-            )
-        elif isinstance(bg_value, int):
-            volume.id = volume.runtime.core.volume_i_from_tiles_device(
-                volume.device.context,
-                ctypes.c_void_p(tile_points.ptr),
-                tile_points.shape[0],
-                transform_buf,
-                translation_buf,
-                in_world_space,
-                bg_value,
-            )
         else:
-            volume.id = volume.runtime.core.volume_f_from_tiles_device(
+            # normalize background value type
+            grid_type = type_to_warp(type(bg_value))
+            if not (is_value(bg_value) or type_is_vector(grid_type)) and (
+                hasattr(bg_value, "__len__") and is_value(bg_value[0])
+            ):
+                # non-warp vectors are considered float, for backward compatibility
+                grid_type = vector(len(bg_value), dtype=float)
+
+            # look for corresponding nvdb type
+            try:
+                nvdb_type = next(
+                    typ
+                    for typ in Volume._supported_allocation_types
+                    if types_equal(grid_type, Volume._nvdb_type_to_dtype[typ])
+                )
+            except StopIteration as err:
+                raise TypeError(
+                    f"Unsupported bg_value type for volume allocation {type_repr(grid_type)}. Supported volume types are {', '.join(Volume._supported_allocation_types)}."
+                ) from err
+
+            # cast to ctype
+            # wrap scalar values in length-1 vectors to handle specific ctype conversion
+            if not type_is_vector(grid_type):
+                grid_type = vector(length=1, dtype=grid_type)
+
+            cvalue = grid_type(bg_value)
+            cvalue_ptr = ctypes.pointer(cvalue)
+            cvalue_size = ctypes.sizeof(cvalue)
+            cvalue_type = nvdb_type.encode("ascii")
+
+            volume.id = volume.runtime.core.volume_from_tiles_device(
                 volume.device.context,
                 ctypes.c_void_p(tile_points.ptr),
                 tile_points.shape[0],
                 transform_buf,
                 translation_buf,
                 in_world_space,
-                float(bg_value),
+                cvalue_ptr,
+                cvalue_size,
+                cvalue_type,
             )
 
         if volume.id == 0:
@@ -4606,6 +4707,8 @@ def matmul(
 ):
     """Computes a generic matrix-matrix multiplication (GEMM) of the form: `d = alpha * (a @ b) + beta * c`.
 
+    .. versionremoved:: 1.7
+
     .. deprecated:: 1.6
         Use :doc:`tile primitives </modules/tiles>` instead.
 
@@ -4619,80 +4722,8 @@ def matmul(
         allow_tf32x3_arith (bool): whether to use CUTLASS's 3xTF32 GEMMs, which enable accuracy similar to FP32
                                    while using Tensor Cores
     """
-    from warp.context import runtime
 
-    warp.utils.warn(
-        "wp.matmul() is deprecated and will be removed in a\nfuture version. Use tile primitives instead.",
-        category=DeprecationWarning,
-        stacklevel=2,
-    )
-
-    device = a.device
-
-    if b.device != device or c.device != device or d.device != device:
-        raise RuntimeError("Matrices A, B, C, and D must all be on the same device as the runtime device.")
-
-    if a.dtype != b.dtype or a.dtype != c.dtype or a.dtype != d.dtype:
-        raise RuntimeError(
-            "wp.matmul currently only supports operation between {A, B, C, D} matrices of the same type."
-        )
-
-    if (
-        (not a.is_contiguous and not a.is_transposed)
-        or (not b.is_contiguous and not b.is_transposed)
-        or (not c.is_contiguous)
-        or (not d.is_contiguous)
-    ):
-        raise RuntimeError(
-            "wp.matmul is only valid for contiguous arrays, with the exception that A and/or B may be transposed."
-        )
-
-    m = a.shape[0]
-    n = b.shape[1]
-    k = a.shape[1]
-    if b.shape != (k, n) or c.shape != (m, n) or d.shape != (m, n):
-        raise RuntimeError(
-            "Invalid shapes for matrices: A = {} B = {} C = {} D = {}".format(a.shape, b.shape, c.shape, d.shape)
-        )
-
-    if runtime.tape:
-        runtime.tape.record_func(
-            backward=lambda: adj_matmul(a, b, c, a.grad, b.grad, c.grad, d.grad, alpha, beta, allow_tf32x3_arith),
-            arrays=[a, b, c, d],
-        )
-        if warp.config.verify_autograd_array_access:
-            d.mark_write()
-            a.mark_read()
-            b.mark_read()
-            c.mark_read()
-
-    # cpu fallback if no cuda devices found
-    if device == "cpu":
-        np_dtype = warp_type_to_np_dtype[a.dtype]
-        d.assign(alpha * np.matmul(a.numpy(), b.numpy(), dtype=np_dtype) + beta * c.numpy())
-        return
-
-    cc = device.arch
-    ret = runtime.core.cutlass_gemm(
-        device.context,
-        cc,
-        m,
-        n,
-        k,
-        type_typestr(a.dtype).encode(),
-        ctypes.c_void_p(a.ptr),
-        ctypes.c_void_p(b.ptr),
-        ctypes.c_void_p(c.ptr),
-        ctypes.c_void_p(d.ptr),
-        alpha,
-        beta,
-        not a.is_transposed,
-        not b.is_transposed,
-        allow_tf32x3_arith,
-        1,
-    )
-    if not ret:
-        raise RuntimeError("matmul failed.")
+    raise RuntimeError("This function has been removed. Use tile primitives instead.")
 
 
 def adj_matmul(
@@ -4724,171 +4755,8 @@ def adj_matmul(
         allow_tf32x3_arith (bool): whether to use CUTLASS's 3xTF32 GEMMs, which enable accuracy similar to FP32
                                    while using Tensor Cores
     """
-    from warp.context import runtime
-
-    device = a.device
-
-    if (
-        b.device != device
-        or c.device != device
-        or adj_a.device != device
-        or adj_b.device != device
-        or adj_c.device != device
-        or adj_d.device != device
-    ):
-        raise RuntimeError(
-            "Matrices A, B, C, D, and their adjoints must all be on the same device as the runtime device."
-        )
-
-    if (
-        a.dtype != b.dtype
-        or a.dtype != c.dtype
-        or a.dtype != adj_a.dtype
-        or a.dtype != adj_b.dtype
-        or a.dtype != adj_c.dtype
-        or a.dtype != adj_d.dtype
-    ):
-        raise RuntimeError(
-            "wp.adj_matmul currently only supports operation between {A, B, C, adj_D, adj_A, adj_B, adj_C} matrices of the same type."
-        )
 
-    if (
-        (not a.is_contiguous and not a.is_transposed)
-        or (not b.is_contiguous and not b.is_transposed)
-        or (not c.is_contiguous)
-        or (not adj_a.is_contiguous and not adj_a.is_transposed)
-        or (not adj_b.is_contiguous and not adj_b.is_transposed)
-        or (not adj_c.is_contiguous)
-        or (not adj_d.is_contiguous)
-    ):
-        raise RuntimeError(
-            "wp.matmul is only valid for contiguous arrays, with the exception that A and/or B and their associated adjoints may be transposed."
-        )
-
-    m = a.shape[0]
-    n = b.shape[1]
-    k = a.shape[1]
-    if (
-        a.shape != (m, k)
-        or b.shape != (k, n)
-        or c.shape != (m, n)
-        or adj_d.shape != (m, n)
-        or adj_a.shape != (m, k)
-        or adj_b.shape != (k, n)
-        or adj_c.shape != (m, n)
-    ):
-        raise RuntimeError(
-            "Invalid shapes for matrices: A = {} B = {} C = {} adj_D = {} adj_A = {} adj_B = {} adj_C = {}".format(
-                a.shape, b.shape, c.shape, adj_d.shape, adj_a.shape, adj_b.shape, adj_c.shape
-            )
-        )
-
-    # cpu fallback if no cuda devices found
-    if device == "cpu":
-        np_dtype = warp_type_to_np_dtype[a.dtype]
-        adj_a.assign(alpha * np.matmul(adj_d.numpy(), b.numpy().transpose(), dtype=np_dtype) + adj_a.numpy())
-        adj_b.assign(alpha * np.matmul(a.numpy().transpose(), adj_d.numpy(), dtype=np_dtype) + adj_b.numpy())
-        adj_c.assign(beta * adj_d.numpy() + adj_c.numpy())
-        return
-
-    cc = device.arch
-
-    # adj_a
-    if not a.is_transposed:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            m,
-            k,
-            n,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(adj_d.ptr),
-            ctypes.c_void_p(b.ptr),
-            ctypes.c_void_p(adj_a.ptr),
-            ctypes.c_void_p(adj_a.ptr),
-            alpha,
-            1.0,
-            True,
-            b.is_transposed,
-            allow_tf32x3_arith,
-            1,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-    else:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            k,
-            m,
-            n,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(b.ptr),
-            ctypes.c_void_p(adj_d.ptr),
-            ctypes.c_void_p(adj_a.ptr),
-            ctypes.c_void_p(adj_a.ptr),
-            alpha,
-            1.0,
-            not b.is_transposed,
-            False,
-            allow_tf32x3_arith,
-            1,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-
-    # adj_b
-    if not b.is_transposed:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            k,
-            n,
-            m,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(a.ptr),
-            ctypes.c_void_p(adj_d.ptr),
-            ctypes.c_void_p(adj_b.ptr),
-            ctypes.c_void_p(adj_b.ptr),
-            alpha,
-            1.0,
-            a.is_transposed,
-            True,
-            allow_tf32x3_arith,
-            1,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-    else:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            n,
-            k,
-            m,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(adj_d.ptr),
-            ctypes.c_void_p(a.ptr),
-            ctypes.c_void_p(adj_b.ptr),
-            ctypes.c_void_p(adj_b.ptr),
-            alpha,
-            1.0,
-            False,
-            not a.is_transposed,
-            allow_tf32x3_arith,
-            1,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-
-    # adj_c
-    warp.launch(
-        kernel=warp.utils.add_kernel_2d,
-        dim=adj_c.shape,
-        inputs=[adj_c, adj_d, adj_d.dtype(beta)],
-        device=device,
-        record_tape=False,
-    )
+    raise RuntimeError("This function has been removed. Use tile primitives instead.")
 
 
 def batched_matmul(
@@ -4902,6 +4770,8 @@ def batched_matmul(
 ):
     """Computes a batched generic matrix-matrix multiplication (GEMM) of the form: `d = alpha * (a @ b) + beta * c`.
 
+    .. versionremoved:: 1.7
+
     .. deprecated:: 1.6
         Use :doc:`tile primitives </modules/tiles>` instead.
 
@@ -4915,107 +4785,8 @@ def batched_matmul(
         allow_tf32x3_arith (bool): whether to use CUTLASS's 3xTF32 GEMMs, which enable accuracy similar to FP32
                                    while using Tensor Cores
     """
-    from warp.context import runtime
-
-    device = a.device
 
-    if b.device != device or c.device != device or d.device != device:
-        raise RuntimeError("Matrices A, B, C, and D must all be on the same device as the runtime device.")
-
-    if a.dtype != b.dtype or a.dtype != c.dtype or a.dtype != d.dtype:
-        raise RuntimeError(
-            "wp.batched_matmul currently only supports operation between {A, B, C, D} matrices of the same type."
-        )
-
-    if (
-        (not a.is_contiguous and not a.is_transposed)
-        or (not b.is_contiguous and not b.is_transposed)
-        or (not c.is_contiguous)
-        or (not d.is_contiguous)
-    ):
-        raise RuntimeError(
-            "wp.matmul is only valid for contiguous arrays, with the exception that A and/or B may be transposed."
-        )
-
-    m = a.shape[1]
-    n = b.shape[2]
-    k = a.shape[2]
-    batch_count = a.shape[0]
-    if b.shape != (batch_count, k, n) or c.shape != (batch_count, m, n) or d.shape != (batch_count, m, n):
-        raise RuntimeError(
-            "Invalid shapes for matrices: A = {} B = {} C = {} D = {}".format(a.shape, b.shape, c.shape, d.shape)
-        )
-
-    if runtime.tape:
-        runtime.tape.record_func(
-            backward=lambda: adj_batched_matmul(
-                a, b, c, a.grad, b.grad, c.grad, d.grad, alpha, beta, allow_tf32x3_arith
-            ),
-            arrays=[a, b, c, d],
-        )
-        if warp.config.verify_autograd_array_access:
-            d.mark_write()
-            a.mark_read()
-            b.mark_read()
-            c.mark_read()
-
-    # cpu fallback if no cuda devices found
-    if device == "cpu":
-        np_dtype = warp_type_to_np_dtype[a.dtype]
-        d.assign(alpha * np.matmul(a.numpy(), b.numpy(), dtype=np_dtype) + beta * c.numpy())
-        return
-
-    # handle case in which batch_count exceeds max_batch_count, which is a CUDA array size maximum
-    max_batch_count = 65535
-    iters = int(batch_count / max_batch_count)
-    remainder = batch_count % max_batch_count
-
-    cc = device.arch
-    for i in range(iters):
-        idx_start = i * max_batch_count
-        idx_end = (i + 1) * max_batch_count if i < iters - 1 else batch_count
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            m,
-            n,
-            k,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(a[idx_start:idx_end, :, :].ptr),
-            ctypes.c_void_p(b[idx_start:idx_end, :, :].ptr),
-            ctypes.c_void_p(c[idx_start:idx_end, :, :].ptr),
-            ctypes.c_void_p(d[idx_start:idx_end, :, :].ptr),
-            alpha,
-            beta,
-            not a.is_transposed,
-            not b.is_transposed,
-            allow_tf32x3_arith,
-            max_batch_count,
-        )
-        if not ret:
-            raise RuntimeError("Batched matmul failed.")
-
-    idx_start = iters * max_batch_count
-    ret = runtime.core.cutlass_gemm(
-        device.context,
-        cc,
-        m,
-        n,
-        k,
-        type_typestr(a.dtype).encode(),
-        ctypes.c_void_p(a[idx_start:, :, :].ptr),
-        ctypes.c_void_p(b[idx_start:, :, :].ptr),
-        ctypes.c_void_p(c[idx_start:, :, :].ptr),
-        ctypes.c_void_p(d[idx_start:, :, :].ptr),
-        alpha,
-        beta,
-        not a.is_transposed,
-        not b.is_transposed,
-        allow_tf32x3_arith,
-        remainder,
-    )
-    if not ret:
-        raise RuntimeError("Batched matmul failed.")
+    raise RuntimeError("This function has been removed. Use tile primitives instead.")
 
 
 def adj_batched_matmul(
@@ -5045,270 +4816,8 @@ def adj_batched_matmul(
         allow_tf32x3_arith (bool): whether to use CUTLASS's 3xTF32 GEMMs, which enable accuracy similar to FP32
                                    while using Tensor Cores
     """
-    from warp.context import runtime
-
-    device = a.device
-
-    if (
-        b.device != device
-        or c.device != device
-        or adj_a.device != device
-        or adj_b.device != device
-        or adj_c.device != device
-        or adj_d.device != device
-    ):
-        raise RuntimeError(
-            "Matrices A, B, C, D, and their adjoints must all be on the same device as the runtime device."
-        )
-
-    if (
-        a.dtype != b.dtype
-        or a.dtype != c.dtype
-        or a.dtype != adj_a.dtype
-        or a.dtype != adj_b.dtype
-        or a.dtype != adj_c.dtype
-        or a.dtype != adj_d.dtype
-    ):
-        raise RuntimeError(
-            "wp.adj_batched_matmul currently only supports operation between {A, B, C, adj_D, adj_A, adj_B, adj_C} matrices of the same type."
-        )
-
-    m = a.shape[1]
-    n = b.shape[2]
-    k = a.shape[2]
-    batch_count = a.shape[0]
-    if (
-        b.shape != (batch_count, k, n)
-        or c.shape != (batch_count, m, n)
-        or adj_d.shape != (batch_count, m, n)
-        or adj_a.shape != (batch_count, m, k)
-        or adj_b.shape != (batch_count, k, n)
-        or adj_c.shape != (batch_count, m, n)
-    ):
-        raise RuntimeError(
-            "Invalid shapes for matrices: A = {} B = {} C = {} adj_D = {} adj_A = {} adj_B = {} adj_C = {}".format(
-                a.shape, b.shape, c.shape, adj_d.shape, adj_a.shape, adj_b.shape, adj_c.shape
-            )
-        )
 
-    if (
-        (not a.is_contiguous and not a.is_transposed)
-        or (not b.is_contiguous and not b.is_transposed)
-        or (not c.is_contiguous)
-        or (not adj_a.is_contiguous and not adj_a.is_transposed)
-        or (not adj_b.is_contiguous and not adj_b.is_transposed)
-        or (not adj_c.is_contiguous)
-        or (not adj_d.is_contiguous)
-    ):
-        raise RuntimeError(
-            "wp.matmul is only valid for contiguous arrays, with the exception that A and/or B and their associated adjoints may be transposed."
-        )
-
-    # cpu fallback if no cuda devices found
-    if device == "cpu":
-        np_dtype = warp_type_to_np_dtype[a.dtype]
-        adj_a.assign(alpha * np.matmul(adj_d.numpy(), b.numpy().transpose((0, 2, 1)), dtype=np_dtype) + adj_a.numpy())
-        adj_b.assign(alpha * np.matmul(a.numpy().transpose((0, 2, 1)), adj_d.numpy(), dtype=np_dtype) + adj_b.numpy())
-        adj_c.assign(beta * adj_d.numpy() + adj_c.numpy())
-        return
-
-    # handle case in which batch_count exceeds max_batch_count, which is a CUDA array size maximum
-    max_batch_count = 65535
-    iters = int(batch_count / max_batch_count)
-    remainder = batch_count % max_batch_count
-
-    cc = device.arch
-
-    for i in range(iters):
-        idx_start = i * max_batch_count
-        idx_end = (i + 1) * max_batch_count if i < iters - 1 else batch_count
-
-        # adj_a
-        if not a.is_transposed:
-            ret = runtime.core.cutlass_gemm(
-                device.context,
-                cc,
-                m,
-                k,
-                n,
-                type_typestr(a.dtype).encode(),
-                ctypes.c_void_p(adj_d[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(b[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_a[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_a[idx_start:idx_end, :, :].ptr),
-                alpha,
-                1.0,
-                True,
-                b.is_transposed,
-                allow_tf32x3_arith,
-                max_batch_count,
-            )
-            if not ret:
-                raise RuntimeError("adj_matmul failed.")
-        else:
-            ret = runtime.core.cutlass_gemm(
-                device.context,
-                cc,
-                k,
-                m,
-                n,
-                type_typestr(a.dtype).encode(),
-                ctypes.c_void_p(b[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_d[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_a[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_a[idx_start:idx_end, :, :].ptr),
-                alpha,
-                1.0,
-                not b.is_transposed,
-                False,
-                allow_tf32x3_arith,
-                max_batch_count,
-            )
-            if not ret:
-                raise RuntimeError("adj_matmul failed.")
-
-        # adj_b
-        if not b.is_transposed:
-            ret = runtime.core.cutlass_gemm(
-                device.context,
-                cc,
-                k,
-                n,
-                m,
-                type_typestr(a.dtype).encode(),
-                ctypes.c_void_p(a[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_d[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_b[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_b[idx_start:idx_end, :, :].ptr),
-                alpha,
-                1.0,
-                a.is_transposed,
-                True,
-                allow_tf32x3_arith,
-                max_batch_count,
-            )
-            if not ret:
-                raise RuntimeError("adj_matmul failed.")
-        else:
-            ret = runtime.core.cutlass_gemm(
-                device.context,
-                cc,
-                n,
-                k,
-                m,
-                type_typestr(a.dtype).encode(),
-                ctypes.c_void_p(adj_d[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(a[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_b[idx_start:idx_end, :, :].ptr),
-                ctypes.c_void_p(adj_b[idx_start:idx_end, :, :].ptr),
-                alpha,
-                1.0,
-                False,
-                not a.is_transposed,
-                allow_tf32x3_arith,
-                max_batch_count,
-            )
-            if not ret:
-                raise RuntimeError("adj_matmul failed.")
-
-    idx_start = iters * max_batch_count
-
-    # adj_a
-    if not a.is_transposed:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            m,
-            k,
-            n,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(adj_d[idx_start:, :, :].ptr),
-            ctypes.c_void_p(b[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_a[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_a[idx_start:, :, :].ptr),
-            alpha,
-            1.0,
-            True,
-            b.is_transposed,
-            allow_tf32x3_arith,
-            remainder,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-    else:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            k,
-            m,
-            n,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(b[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_d[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_a[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_a[idx_start:, :, :].ptr),
-            alpha,
-            1.0,
-            not b.is_transposed,
-            False,
-            allow_tf32x3_arith,
-            remainder,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-
-    # adj_b
-    if not b.is_transposed:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            k,
-            n,
-            m,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(a[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_d[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_b[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_b[idx_start:, :, :].ptr),
-            alpha,
-            1.0,
-            a.is_transposed,
-            True,
-            allow_tf32x3_arith,
-            remainder,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-    else:
-        ret = runtime.core.cutlass_gemm(
-            device.context,
-            cc,
-            n,
-            k,
-            m,
-            type_typestr(a.dtype).encode(),
-            ctypes.c_void_p(adj_d[idx_start:, :, :].ptr),
-            ctypes.c_void_p(a[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_b[idx_start:, :, :].ptr),
-            ctypes.c_void_p(adj_b[idx_start:, :, :].ptr),
-            alpha,
-            1.0,
-            False,
-            not a.is_transposed,
-            allow_tf32x3_arith,
-            remainder,
-        )
-        if not ret:
-            raise RuntimeError("adj_matmul failed.")
-
-    # adj_c
-    warp.launch(
-        kernel=warp.utils.add_kernel_3d,
-        dim=adj_c.shape,
-        inputs=[adj_c, adj_d, adj_d.dtype(beta)],
-        device=device,
-        record_tape=False,
-    )
+    raise RuntimeError("This function has been removed. Use tile primitives instead.")
 
 
 class HashGrid:
@@ -5691,7 +5200,7 @@ simple_type_codes = {
 }
 
 
-def get_type_code(arg_type):
+def get_type_code(arg_type: type) -> str:
     if arg_type == Any:
         # special case for generics
         # note: since Python 3.11 Any is a type, so we check for it first
@@ -5755,8 +5264,8 @@ def get_type_code(arg_type):
         raise TypeError(f"Unrecognized type '{arg_type}'")
 
 
-def get_signature(arg_types, func_name=None, arg_names=None):
-    type_codes = []
+def get_signature(arg_types: List[type], func_name: Optional[str] = None, arg_names: Optional[List[str]] = None) -> str:
+    type_codes: List[str] = []
     for i, arg_type in enumerate(arg_types):
         try:
             type_codes.append(get_type_code(arg_type))