warp-lang 1.9.1__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/utils.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,1668 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
+# isort: skip_file
 
-import cProfile
-import ctypes
-import os
-import sys
-import time
-import warnings
-from types import ModuleType
-from typing import Any, Callable
+from warp._src.utils import array_cast as array_cast
+from warp._src.utils import segmented_sort_pairs as segmented_sort_pairs
+from warp._src.utils import warn as warn
 
-import numpy as np
 
-import warp as wp
-import warp.context
-import warp.types
-from warp.context import Devicelike
-from warp.types import Array, DType, type_repr, types_equal
+# TODO: Remove after cleaning up the public API.
 
-warnings_seen = set()
+from warp._src import utils as _utils
 
 
-def warp_showwarning(message, category, filename, lineno, file=None, line=None):
-    """Version of warnings.showwarning that always prints to sys.stdout."""
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
 
-    if warp.config.verbose_warnings:
-        s = f"Warp {category.__name__}: {message} ({filename}:{lineno})\n"
-
-        if line is None:
-            try:
-                import linecache
-
-                line = linecache.getline(filename, lineno)
-            except Exception:
-                # When a warning is logged during Python shutdown, linecache
-                # and the import machinery don't work anymore
-                line = None
-                linecache = None
-
-        if line:
-            line = line.strip()
-            s += f"  {line}\n"
-    else:
-        # simple warning
-        s = f"Warp {category.__name__}: {message}\n"
-
-    sys.stdout.write(s)
-
-
-def warn(message, category=None, stacklevel=1):
-    if (category, message) in warnings_seen:
-        return
-
-    with warnings.catch_warnings():
-        warnings.simplefilter("default")  # Change the filter in this process
-        warnings.showwarning = warp_showwarning
-        warnings.warn(
-            message,
-            category,
-            stacklevel=stacklevel + 1,  # Increment stacklevel by 1 since we are in a wrapper
-        )
-
-    if category is DeprecationWarning:
-        warnings_seen.add((category, message))
-
-
-# expand a 7-vec to a tuple of arrays
-def transform_expand(t):
-    return wp.transform(np.array(t[0:3]), np.array(t[3:7]))
-
-
-@wp.func
-def quat_between_vectors(a: wp.vec3, b: wp.vec3) -> wp.quat:
-    """
-    Compute the quaternion that rotates vector a to vector b
-    """
-    a = wp.normalize(a)
-    b = wp.normalize(b)
-    c = wp.cross(a, b)
-    d = wp.dot(a, b)
-    q = wp.quat(c[0], c[1], c[2], 1.0 + d)
-    return wp.normalize(q)
-
-
-def array_scan(in_array, out_array, inclusive=True):
-    """Perform a scan (prefix sum) operation on an array.
-
-    This function computes the inclusive or exclusive scan of the input array and stores the result in the output array.
-    The scan operation computes a running sum of elements in the array.
-
-    Args:
-        in_array (wp.array): Input array to scan. Must be of type int32 or float32.
-        out_array (wp.array): Output array to store scan results. Must match input array type and size.
-        inclusive (bool, optional): If True, performs an inclusive scan (includes current element in sum).
-                                  If False, performs an exclusive scan (excludes current element). Defaults to True.
-
-    Raises:
-        RuntimeError: If array storage devices don't match, if storage size is insufficient, or if data types are unsupported.
-    """
-
-    if in_array.device != out_array.device:
-        raise RuntimeError(f"In and out array storage devices do not match ({in_array.device} vs {out_array.device})")
-
-    if in_array.size != out_array.size:
-        raise RuntimeError(f"In and out array storage sizes do not match ({in_array.size} vs {out_array.size})")
-
-    if not types_equal(in_array.dtype, out_array.dtype):
-        raise RuntimeError(
-            f"In and out array data types do not match ({type_repr(in_array.dtype)} vs {type_repr(out_array.dtype)})"
-        )
-
-    if in_array.size == 0:
-        return
-
-    from warp.context import runtime
-
-    if in_array.device.is_cpu:
-        if in_array.dtype == wp.int32:
-            runtime.core.wp_array_scan_int_host(in_array.ptr, out_array.ptr, in_array.size, inclusive)
-        elif in_array.dtype == wp.float32:
-            runtime.core.wp_array_scan_float_host(in_array.ptr, out_array.ptr, in_array.size, inclusive)
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(in_array.dtype)}")
-    elif in_array.device.is_cuda:
-        if in_array.dtype == wp.int32:
-            runtime.core.wp_array_scan_int_device(in_array.ptr, out_array.ptr, in_array.size, inclusive)
-        elif in_array.dtype == wp.float32:
-            runtime.core.wp_array_scan_float_device(in_array.ptr, out_array.ptr, in_array.size, inclusive)
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(in_array.dtype)}")
-
-
-def radix_sort_pairs(keys, values, count: int):
-    """Sort key-value pairs using radix sort.
-
-    This function sorts pairs of arrays based on the keys array, maintaining the key-value
-    relationship. The sort is stable and operates in linear time.
-    The `keys` and `values` arrays must be large enough to accommodate 2*`count` elements.
-
-    Args:
-        keys (wp.array): Array of keys to sort. Must be of type int32, float32, or int64.
-        values (wp.array): Array of values to sort along with keys. Must be of type int32.
-        count (int): Number of elements to sort.
-
-    Raises:
-        RuntimeError: If array storage devices don't match, if storage size is insufficient, or if data types are unsupported.
-    """
-    if keys.device != values.device:
-        raise RuntimeError(f"Keys and values array storage devices do not match ({keys.device} vs {values.device})")
-
-    if count == 0:
-        return
-
-    if keys.size < 2 * count or values.size < 2 * count:
-        raise RuntimeError("Keys and values array storage must be large enough to contain 2*count elements")
-
-    from warp.context import runtime
-
-    if keys.device.is_cpu:
-        if keys.dtype == wp.int32 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_int_host(keys.ptr, values.ptr, count)
-        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_float_host(keys.ptr, values.ptr, count)
-        elif keys.dtype == wp.int64 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_int64_host(keys.ptr, values.ptr, count)
-        else:
-            raise RuntimeError(
-                f"Unsupported keys and values data types: {type_repr(keys.dtype)}, {type_repr(values.dtype)}"
-            )
-    elif keys.device.is_cuda:
-        if keys.dtype == wp.int32 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_int_device(keys.ptr, values.ptr, count)
-        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_float_device(keys.ptr, values.ptr, count)
-        elif keys.dtype == wp.int64 and values.dtype == wp.int32:
-            runtime.core.wp_radix_sort_pairs_int64_device(keys.ptr, values.ptr, count)
-        else:
-            raise RuntimeError(
-                f"Unsupported keys and values data types: {type_repr(keys.dtype)}, {type_repr(values.dtype)}"
-            )
-
-
-def segmented_sort_pairs(
-    keys,
-    values,
-    count: int,
-    segment_start_indices: wp.array(dtype=wp.int32),
-    segment_end_indices: wp.array(dtype=wp.int32) = None,
-):
-    """Sort key-value pairs within segments.
-
-    This function performs a segmented sort of key-value pairs, where the sorting is done independently within each segment.
-    The segments are defined by their start and optionally end indices.
-    The `keys` and `values` arrays must be large enough to accommodate 2*`count` elements.
-
-    Args:
-        keys: Array of keys to sort. Must be of type int32 or float32.
-        values: Array of values to sort along with keys. Must be of type int32.
-        count: Number of elements to sort.
-        segment_start_indices: Array containing start index of each segment. Must be of type int32.
-            If segment_end_indices is None, this array must have length at least num_segments + 1,
-            and segment_end_indices will be inferred as segment_start_indices[1:].
-            If segment_end_indices is provided, this array must have length at least num_segments.
-        segment_end_indices: Optional array containing end index of each segment. Must be of type int32 if provided.
-            If None, segment_end_indices will be inferred from segment_start_indices[1:].
-            If provided, must have length at least num_segments.
-
-    Raises:
-        RuntimeError: If array storage devices don't match, if storage size is insufficient,
-                     if segment_start_indices is not of type int32, or if data types are unsupported.
-    """
-    if keys.device != values.device:
-        raise RuntimeError(f"Array storage devices do not match ({keys.device} vs {values.device})")
-
-    if count == 0:
-        return
-
-    if keys.size < 2 * count or values.size < 2 * count:
-        raise RuntimeError("Array storage must be large enough to contain 2*count elements")
-
-    from warp.context import runtime
-
-    if segment_start_indices.dtype != wp.int32:
-        raise RuntimeError("segment_start_indices array must be of type int32")
-
-    # Handle case where segment_end_indices is not provided
-    if segment_end_indices is None:
-        num_segments = max(0, segment_start_indices.size - 1)
-
-        segment_end_indices = segment_start_indices[1:]
-        segment_end_indices_ptr = segment_end_indices.ptr
-        segment_start_indices_ptr = segment_start_indices.ptr
-    else:
-        if segment_end_indices.dtype != wp.int32:
-            raise RuntimeError("segment_end_indices array must be of type int32")
-
-        num_segments = segment_start_indices.size
-
-        segment_end_indices_ptr = segment_end_indices.ptr
-        segment_start_indices_ptr = segment_start_indices.ptr
-
-    if keys.device.is_cpu:
-        if keys.dtype == wp.int32 and values.dtype == wp.int32:
-            runtime.core.wp_segmented_sort_pairs_int_host(
-                keys.ptr,
-                values.ptr,
-                count,
-                segment_start_indices_ptr,
-                segment_end_indices_ptr,
-                num_segments,
-            )
-        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
-            runtime.core.wp_segmented_sort_pairs_float_host(
-                keys.ptr,
-                values.ptr,
-                count,
-                segment_start_indices_ptr,
-                segment_end_indices_ptr,
-                num_segments,
-            )
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(keys.dtype)}")
-    elif keys.device.is_cuda:
-        if keys.dtype == wp.int32 and values.dtype == wp.int32:
-            runtime.core.wp_segmented_sort_pairs_int_device(
-                keys.ptr,
-                values.ptr,
-                count,
-                segment_start_indices_ptr,
-                segment_end_indices_ptr,
-                num_segments,
-            )
-        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
-            runtime.core.wp_segmented_sort_pairs_float_device(
-                keys.ptr,
-                values.ptr,
-                count,
-                segment_start_indices_ptr,
-                segment_end_indices_ptr,
-                num_segments,
-            )
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(keys.dtype)}")
-
-
-def runlength_encode(values, run_values, run_lengths, run_count=None, value_count=None):
-    """Perform run-length encoding on an array.
-
-    This function compresses an array by replacing consecutive identical values with a single value
-    and its count. For example, [1,1,1,2,2,3] becomes values=[1,2,3] and lengths=[3,2,1].
-
-    Args:
-        values (wp.array): Input array to encode. Must be of type int32.
-        run_values (wp.array): Output array to store unique values. Must be at least value_count in size.
-        run_lengths (wp.array): Output array to store run lengths. Must be at least value_count in size.
-        run_count (wp.array, optional): Optional output array to store the number of runs.
-                                       If None, returns the count as an integer.
-        value_count (int, optional): Number of values to process. If None, processes entire array.
-
-    Returns:
-        int or wp.array: Number of runs if run_count is None, otherwise returns run_count array.
-
-    Raises:
-        RuntimeError: If array storage devices don't match, if storage size is insufficient, or if data types are unsupported.
-    """
-    if run_values.device != values.device or run_lengths.device != values.device:
-        raise RuntimeError("run_values, run_lengths and values storage devices do not match")
-
-    if value_count is None:
-        value_count = values.size
-
-    if run_values.size < value_count or run_lengths.size < value_count:
-        raise RuntimeError(f"Output array storage sizes must be at least equal to value_count ({value_count})")
-
-    if not types_equal(values.dtype, run_values.dtype):
-        raise RuntimeError(
-            f"values and run_values data types do not match ({type_repr(values.dtype)} vs {type_repr(run_values.dtype)})"
-        )
-
-    if run_lengths.dtype != wp.int32:
-        raise RuntimeError("run_lengths array must be of type int32")
-
-    # User can provide a device output array for storing the number of runs
-    # For convenience, if no such array is provided, number of runs is returned on host
-    if run_count is None:
-        if value_count == 0:
-            return 0
-        run_count = wp.empty(shape=(1,), dtype=int, device=values.device)
-        host_return = True
-    else:
-        if run_count.device != values.device:
-            raise RuntimeError("run_count storage device does not match other arrays")
-        if run_count.dtype != wp.int32:
-            raise RuntimeError("run_count array must be of type int32")
-        if value_count == 0:
-            run_count.zero_()
-            return run_count
-        host_return = False
-
-    from warp.context import runtime
-
-    if values.device.is_cpu:
-        if values.dtype == wp.int32:
-            runtime.core.wp_runlength_encode_int_host(
-                values.ptr, run_values.ptr, run_lengths.ptr, run_count.ptr, value_count
-            )
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(values.dtype)}")
-    elif values.device.is_cuda:
-        if values.dtype == wp.int32:
-            runtime.core.wp_runlength_encode_int_device(
-                values.ptr, run_values.ptr, run_lengths.ptr, run_count.ptr, value_count
-            )
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(values.dtype)}")
-
-    if host_return:
-        return int(run_count.numpy()[0])
-    return run_count
-
-
-def array_sum(values, out=None, value_count=None, axis=None):
-    """Compute the sum of array elements.
-
-    This function computes the sum of array elements, optionally along a specified axis.
-    The operation can be performed on the entire array or along a specific dimension.
-
-    Args:
-        values (wp.array): Input array to sum. Must be of type float32 or float64.
-        out (wp.array, optional): Output array to store results. If None, a new array is created.
-        value_count (int, optional): Number of elements to process. If None, processes entire array.
-        axis (int, optional): Axis along which to compute sum. If None, computes sum of all elements.
-
-    Returns:
-        wp.array or float: The sum result. Returns a float if axis is None and out is None,
-                           otherwise returns the output array.
-
-    Raises:
-        RuntimeError: If output array storage device or data type is incompatible with input array.
-    """
-    if value_count is None:
-        if axis is None:
-            value_count = values.size
-        else:
-            value_count = values.shape[axis]
-
-    if axis is None:
-        output_shape = (1,)
-    else:
-
-        def output_dim(ax, dim):
-            return 1 if ax == axis else dim
-
-        output_shape = tuple(output_dim(ax, dim) for ax, dim in enumerate(values.shape))
-
-    type_size = wp.types.type_size(values.dtype)
-    scalar_type = wp.types.type_scalar_type(values.dtype)
-
-    # User can provide a device output array for storing the number of runs
-    # For convenience, if no such array is provided, number of runs is returned on host
-    if out is None:
-        host_return = True
-        out = wp.empty(shape=output_shape, dtype=values.dtype, device=values.device)
-    else:
-        host_return = False
-        if out.device != values.device:
-            raise RuntimeError("out storage device should match values array")
-        if out.dtype != values.dtype:
-            raise RuntimeError(f"out array should have type {values.dtype.__name__}")
-        if out.shape != output_shape:
-            raise RuntimeError(f"out array should have shape {output_shape}")
-
-    if value_count == 0:
-        out.zero_()
-        if axis is None and host_return:
-            return out.numpy()[0]
-        return out
-
-    from warp.context import runtime
-
-    if values.device.is_cpu:
-        if scalar_type == wp.float32:
-            native_func = runtime.core.wp_array_sum_float_host
-        elif scalar_type == wp.float64:
-            native_func = runtime.core.wp_array_sum_double_host
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(values.dtype)}")
-    elif values.device.is_cuda:
-        if scalar_type == wp.float32:
-            native_func = runtime.core.wp_array_sum_float_device
-        elif scalar_type == wp.float64:
-            native_func = runtime.core.wp_array_sum_double_device
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(values.dtype)}")
-
-    if axis is None:
-        stride = wp.types.type_size_in_bytes(values.dtype)
-        native_func(values.ptr, out.ptr, value_count, stride, type_size)
-
-        if host_return:
-            return out.numpy()[0]
-        return out
-
-    stride = values.strides[axis]
-    for idx in np.ndindex(output_shape):
-        out_offset = sum(i * s for i, s in zip(idx, out.strides))
-        val_offset = sum(i * s for i, s in zip(idx, values.strides))
-
-        native_func(
-            values.ptr + val_offset,
-            out.ptr + out_offset,
-            value_count,
-            stride,
-            type_size,
-        )
-
-    return out
-
-
-def array_inner(a, b, out=None, count=None, axis=None):
-    """Compute the inner product of two arrays.
-
-    This function computes the dot product between two arrays, optionally along a specified axis.
-    The operation can be performed on the entire arrays or along a specific dimension.
-
-    Args:
-        a (wp.array): First input array.
-        b (wp.array): Second input array. Must match shape and type of a.
-        out (wp.array, optional): Output array to store results. If None, a new array is created.
-        count (int, optional): Number of elements to process. If None, processes entire arrays.
-        axis (int, optional): Axis along which to compute inner product. If None, computes on flattened arrays.
-
-    Returns:
-        wp.array or float: The inner product result. Returns a float if axis is None and out is None,
-                           otherwise returns the output array.
-
-    Raises:
-        RuntimeError: If array storage devices, sizes, or data types are incompatible.
-    """
-    if a.size != b.size:
-        raise RuntimeError(f"A and b array storage sizes do not match ({a.size} vs {b.size})")
-
-    if a.device != b.device:
-        raise RuntimeError(f"A and b array storage devices do not match ({a.device} vs {b.device})")
-
-    if not types_equal(a.dtype, b.dtype):
-        raise RuntimeError(f"A and b array data types do not match ({type_repr(a.dtype)} vs {type_repr(b.dtype)})")
-
-    if count is None:
-        if axis is None:
-            count = a.size
-        else:
-            count = a.shape[axis]
-
-    if axis is None:
-        output_shape = (1,)
-    else:
-
-        def output_dim(ax, dim):
-            return 1 if ax == axis else dim
-
-        output_shape = tuple(output_dim(ax, dim) for ax, dim in enumerate(a.shape))
-
-    type_size = wp.types.type_size(a.dtype)
-    scalar_type = wp.types.type_scalar_type(a.dtype)
-
-    # User can provide a device output array for storing the number of runs
-    # For convenience, if no such array is provided, number of runs is returned on host
-    if out is None:
-        host_return = True
-        out = wp.empty(shape=output_shape, dtype=scalar_type, device=a.device)
-    else:
-        host_return = False
-        if out.device != a.device:
-            raise RuntimeError("out storage device should match values array")
-        if out.dtype != scalar_type:
-            raise RuntimeError(f"out array should have type {scalar_type.__name__}")
-        if out.shape != output_shape:
-            raise RuntimeError(f"out array should have shape {output_shape}")
-
-    if count == 0:
-        if axis is None and host_return:
-            return 0.0
-        out.zero_()
-        return out
-
-    from warp.context import runtime
-
-    if a.device.is_cpu:
-        if scalar_type == wp.float32:
-            native_func = runtime.core.wp_array_inner_float_host
-        elif scalar_type == wp.float64:
-            native_func = runtime.core.wp_array_inner_double_host
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(a.dtype)}")
-    elif a.device.is_cuda:
-        if scalar_type == wp.float32:
-            native_func = runtime.core.wp_array_inner_float_device
-        elif scalar_type == wp.float64:
-            native_func = runtime.core.wp_array_inner_double_device
-        else:
-            raise RuntimeError(f"Unsupported data type: {type_repr(a.dtype)}")
-
-    if axis is None:
-        stride_a = wp.types.type_size_in_bytes(a.dtype)
-        stride_b = wp.types.type_size_in_bytes(b.dtype)
-        native_func(a.ptr, b.ptr, out.ptr, count, stride_a, stride_b, type_size)
-
-        if host_return:
-            return out.numpy()[0]
-        return out
-
-    stride_a = a.strides[axis]
-    stride_b = b.strides[axis]
-
-    for idx in np.ndindex(output_shape):
-        out_offset = sum(i * s for i, s in zip(idx, out.strides))
-        a_offset = sum(i * s for i, s in zip(idx, a.strides))
-        b_offset = sum(i * s for i, s in zip(idx, b.strides))
-
-        native_func(
-            a.ptr + a_offset,
-            b.ptr + b_offset,
-            out.ptr + out_offset,
-            count,
-            stride_a,
-            stride_b,
-            type_size,
-        )
-
-    return out
-
-
-@wp.kernel
-def _array_cast_kernel(
-    dest: Any,
-    src: Any,
-):
-    i = wp.tid()
-    dest[i] = dest.dtype(src[i])
-
-
-def array_cast(in_array, out_array, count=None):
-    """Cast elements from one array to another array with a different data type.
-
-    This function performs element-wise casting from the input array to the output array.
-    The arrays must have the same number of dimensions and data type shapes. If they don't match,
-    the arrays will be flattened and casting will be performed at the scalar level.
-
-    Args:
-        in_array (wp.array): Input array to cast from.
-        out_array (wp.array): Output array to cast to. Must have the same device as in_array.
-        count (int, optional): Number of elements to process. If None, processes entire array.
-                             For multi-dimensional arrays, partial casting is not supported.
-
-    Raises:
-        RuntimeError: If arrays have different devices or if attempting partial casting
-                     on multi-dimensional arrays.
-
-    Note:
-        If the input and output arrays have the same data type, this function will
-        simply copy the data without any conversion.
-    """
-    if in_array.device != out_array.device:
-        raise RuntimeError(f"Array storage devices do not match ({in_array.device} vs {out_array.device})")
-
-    in_array_data_shape = getattr(in_array.dtype, "_shape_", ())
-    out_array_data_shape = getattr(out_array.dtype, "_shape_", ())
-
-    if in_array.ndim != out_array.ndim or in_array_data_shape != out_array_data_shape:
-        # Number of dimensions or data type shape do not match.
-        # Flatten arrays and do cast at the scalar level
-        in_array = in_array.flatten()
-        out_array = out_array.flatten()
-
-        in_array_data_length = warp.types.type_size(in_array.dtype)
-        out_array_data_length = warp.types.type_size(out_array.dtype)
-        in_array_scalar_type = wp.types.type_scalar_type(in_array.dtype)
-        out_array_scalar_type = wp.types.type_scalar_type(out_array.dtype)
-
-        in_array = wp.array(
-            data=None,
-            ptr=in_array.ptr,
-            capacity=in_array.capacity,
-            device=in_array.device,
-            dtype=in_array_scalar_type,
-            shape=in_array.shape[0] * in_array_data_length,
-        )
-
-        out_array = wp.array(
-            data=None,
-            ptr=out_array.ptr,
-            capacity=out_array.capacity,
-            device=out_array.device,
-            dtype=out_array_scalar_type,
-            shape=out_array.shape[0] * out_array_data_length,
-        )
-
-        if count is not None:
-            count *= in_array_data_length
-
-    if count is None:
-        count = in_array.size
-
-    if in_array.ndim == 1:
-        dim = count
-    elif count < in_array.size:
-        raise RuntimeError("Partial cast is not supported for arrays with more than one dimension")
-    else:
-        dim = in_array.shape
-
-    if in_array.dtype == out_array.dtype:
-        # Same data type, can simply copy
-        wp.copy(dest=out_array, src=in_array, count=count)
-    else:
-        wp.launch(kernel=_array_cast_kernel, dim=dim, inputs=[out_array, in_array], device=out_array.device)
-
-
-def create_warp_function(func: Callable) -> tuple[wp.Function, warp.context.Module]:
-    """Create a Warp function from a Python function.
-
-    Args:
-        func (Callable): A Python function to be converted to a Warp function.
-
-    Returns:
-        wp.Function: A Warp function created from the input function.
-    """
-
-    from .codegen import Adjoint, get_full_arg_spec
-
-    def unique_name(code: str):
-        return "func_" + hex(hash(code))[-8:]
-
-    # Create a Warp function from the input function
-    source = None
-    argspec = get_full_arg_spec(func)
-    key = getattr(func, "__name__", None)
-    if key is None:
-        source, _ = Adjoint.extract_function_source(func)
-        key = unique_name(source)
-    elif key == "<lambda>":
-        body = Adjoint.extract_lambda_source(func, only_body=True)
-        if body is None:
-            raise ValueError("Could not extract lambda source code")
-        key = unique_name(body)
-        source = f"def {key}({', '.join(argspec.args)}):\n  return {body}"
-    else:
-        # use the qualname of the function as the key
-        key = getattr(func, "__qualname__", key)
-        key = key.replace(".", "_").replace(" ", "_").replace("<", "").replace(">", "_")
-
-    module = warp.context.get_module(f"map_{key}")
-    func = wp.Function(
-        func,
-        namespace="",
-        module=module,
-        key=key,
-        source=source,
-        overloaded_annotations=dict.fromkeys(argspec.args, Any),
-    )
-    return func, module
-
-
-def broadcast_shapes(shapes: list[tuple[int]]) -> tuple[int]:
-    """Broadcast a list of shapes to a common shape.
-
-    Following the broadcasting rules of NumPy, two shapes are compatible when:
-    starting from the trailing dimension,
-        1. the two dimensions are equal, or
-        2. one of the dimensions is 1.
-
-    Example:
-        >>> broadcast_shapes([(3, 1, 4), (5, 4)])
-        (3, 5, 4)
-
-    Returns:
-        tuple[int]: The broadcasted shape.
-
-    Raises:
-        ValueError: If the shapes are not broadcastable.
-    """
-    ref = shapes[0]
-    for shape in shapes[1:]:
-        broad = []
-        for j in range(1, max(len(ref), len(shape)) + 1):
-            if j <= len(ref) and j <= len(shape):
-                s = shape[-j]
-                r = ref[-j]
-                if s == r:
-                    broad.append(s)
-                elif s == 1 or r == 1:
-                    broad.append(max(s, r))
-                else:
-                    raise ValueError(f"Shapes {ref} and {shape} are not broadcastable")
-            elif j <= len(ref):
-                broad.append(ref[-j])
-            else:
-                broad.append(shape[-j])
-        ref = tuple(reversed(broad))
-    return ref
-
-
-def map(
-    func: Callable | wp.Function,
-    *inputs: Array[DType] | Any,
-    out: Array[DType] | list[Array[DType]] | None = None,
-    return_kernel: bool = False,
-    block_dim=256,
-    device: Devicelike = None,
-) -> Array[DType] | list[Array[DType]] | wp.Kernel:
-    """
-    Map a function over the elements of one or more arrays.
-
-    You can use a Warp function, a regular Python function, or a lambda expression to map it to a set of arrays.
-
-    .. testcode::
-
-        a = wp.array([1, 2, 3], dtype=wp.float32)
-        b = wp.array([4, 5, 6], dtype=wp.float32)
-        c = wp.array([7, 8, 9], dtype=wp.float32)
-        result = wp.map(lambda x, y, z: x + 2.0 * y - z, a, b, c)
-        print(result)
-
-    .. testoutput::
-
-        [2. 4. 6.]
-
-    Clamp values in an array in place:
-
-    .. testcode::
-
-        xs = wp.array([-1.0, 0.0, 1.0], dtype=wp.float32)
-        wp.map(wp.clamp, xs, -0.5, 0.5, out=xs)
-        print(xs)
-
-    .. testoutput::
-
-        [-0.5  0.   0.5]
-
-    Note that only one of the inputs must be a Warp array. For example, it is possible
-    vectorize the function :func:`warp.transform_point` over a collection of points
-    with a given input transform as follows:
-
-    .. code-block:: python
-
-        tf = wp.transform((1.0, 2.0, 3.0), wp.quat_rpy(0.2, -0.6, 0.1))
-        points = wp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=wp.vec3)
-        transformed = wp.map(wp.transform_point, tf, points)
-
-    Besides regular Warp arrays, other array types, such as the ``indexedarray``, are supported as well:
-
-    .. testcode::
-
-        arr = wp.array(data=np.arange(10, dtype=np.float32))
-        indices = wp.array([1, 3, 5, 7, 9], dtype=int)
-        iarr = wp.indexedarray1d(arr, [indices])
-        out = wp.map(lambda x: x * 10.0, iarr)
-        print(out)
-
-    .. testoutput::
-
-        [10. 30. 50. 70. 90.]
-
-    If multiple arrays are provided, the
-    `NumPy broadcasting rules <https://numpy.org/doc/stable/user/basics.broadcasting.html>`_
-    are applied to determine the shape of the output array.
-    Two shapes are compatible when:
-    starting from the trailing dimension,
-
-    1. the two dimensions are equal, or
-    2. one of the dimensions is 1.
-
-    For example, given arrays of shapes ``(3, 1, 4)`` and ``(5, 4)``, the broadcasted
-    shape is ``(3, 5, 4)``.
-
-    If no array(s) are provided to the ``out`` argument, the output array(s) are created automatically.
-    The data type(s) of the output array(s) are determined by the type of the return value(s) of
-    the function. The ``requires_grad`` flag for an automatically created output array is set to ``True``
-    if any of the input arrays have it set to ``True`` and the respective output array's ``dtype`` is a type that
-    supports differentiation.
-
-    Args:
-        func (Callable | Function): The function to map over the arrays.
-        *inputs (array | Any): The input arrays or values to pass to the function.
-        out (array | list[array] | None): Optional output array(s) to store the result(s). If None, the output array(s) will be created automatically.
-        return_kernel (bool): If True, only return the generated kernel without performing the mapping operation.
-        block_dim (int): The block dimension for the kernel launch.
-        device (Devicelike): The device on which to run the kernel.
-
-    Returns:
-        array | list[array] | Kernel:
-            The resulting array(s) of the mapping. If ``return_kernel`` is True, only returns the kernel used for mapping.
-    """
-
-    import builtins
-
-    from .codegen import Adjoint, Struct, StructInstance
-    from .types import (
-        is_array,
-        type_is_matrix,
-        type_is_quaternion,
-        type_is_transformation,
-        type_is_vector,
-        type_repr,
-        type_to_warp,
-        types_equal,
-    )
-
-    # mapping from struct name to its Python definition
-    referenced_modules: dict[str, ModuleType] = {}
-
-    def type_to_code(wp_type) -> str:
-        """Returns the string representation of a given Warp type."""
-        if is_array(wp_type):
-            return f"warp.array(ndim={wp_type.ndim}, dtype={type_to_code(wp_type.dtype)})"
-        if isinstance(wp_type, Struct):
-            key = f"{wp_type.__module__}.{wp_type.key}"
-            module = sys.modules.get(wp_type.__module__, None)
-            if module is not None:
-                referenced_modules[wp_type.__module__] = module
-            return key
-        if type_is_transformation(wp_type):
-            return f"warp.types.transformation(dtype={type_to_code(wp_type._wp_scalar_type_)})"
-        if type_is_quaternion(wp_type):
-            return f"warp.types.quaternion(dtype={type_to_code(wp_type._wp_scalar_type_)})"
-        if type_is_vector(wp_type):
-            return f"warp.types.vector(length={wp_type._shape_[0]}, dtype={type_to_code(wp_type._wp_scalar_type_)})"
-        if type_is_matrix(wp_type):
-            return f"warp.types.matrix(shape=({wp_type._shape_[0]}, {wp_type._shape_[1]}), dtype={type_to_code(wp_type._wp_scalar_type_)})"
-        if wp_type == builtins.bool:
-            return "bool"
-        if wp_type == builtins.float:
-            return "float"
-        if wp_type == builtins.int:
-            return "int"
-
-        name = getattr(wp_type, "__name__", None)
-        if name is None:
-            return type_repr(wp_type)
-        name = getattr(wp_type, "__qualname__", name)
-        module = getattr(wp_type, "__module__", None)
-        if module is not None:
-            referenced_modules[wp_type.__module__] = module
-        return wp_type.__module__ + "." + name
-
-    def get_warp_type(value):
-        dtype = type(value)
-        if issubclass(dtype, StructInstance):
-            # a struct
-            return value._cls
-        return type_to_warp(dtype)
-
-    # gather the arrays in the inputs
-    array_shapes = [a.shape for a in inputs if is_array(a)]
-    if len(array_shapes) == 0:
-        raise ValueError("map requires at least one warp.array input")
-    # broadcast the shapes of the arrays
-    out_shape = broadcast_shapes(array_shapes)
-
-    module = None
-    out_dtypes = None
-    if isinstance(func, wp.Function):
-        func_name = func.key
-        wp_func = func
-    else:
-        # check if op is a callable function
-        if not callable(func):
-            raise TypeError("func must be a callable function or a warp.Function")
-        wp_func, module = create_warp_function(func)
-        func_name = wp_func.key
-    if module is None:
-        module = warp.context.get_module(f"map_{func_name}")
-
-    arg_names = list(wp_func.input_types.keys())
-
-    if len(inputs) != len(arg_names):
-        raise TypeError(
-            f"Number of input arguments ({len(inputs)}) does not match expected number of function arguments ({len(arg_names)})"
-        )
-
-    # determine output dtype
-    arg_types = {}
-    arg_values = {}
-    for i, arg_name in enumerate(arg_names):
-        if is_array(inputs[i]):
-            # we will pass an element of the array to the function
-            arg_types[arg_name] = inputs[i].dtype
-            if device is None:
-                device = inputs[i].device
-        else:
-            # we pass the input value directly to the function
-            arg_types[arg_name] = get_warp_type(inputs[i])
-    func_or_none = wp_func.get_overload(list(arg_types.values()), {})
-    if func_or_none is None:
-        raise TypeError(
-            f"Function {func_name} does not support the provided argument types {', '.join(type_repr(t) for t in arg_types.values())}"
-        )
-    func = func_or_none
-
-    if func.value_type is not None:
-        out_dtype = func.value_type
-    elif func.value_func is not None:
-        out_dtype = func.value_func(arg_types, arg_values)
-    else:
-        func.build(None)
-        out_dtype = func.value_func(arg_types, arg_values)
-
-    if out_dtype is None:
-        raise TypeError("The provided function must return a value")
-
-    if isinstance(out_dtype, tuple) or isinstance(out_dtype, list):
-        out_dtypes = out_dtype
-    else:
-        out_dtypes = (out_dtype,)
-
-    if out is None:
-        requires_grad = any(getattr(a, "requires_grad", False) for a in inputs if is_array(a))
-        outputs = []
-        for dtype in out_dtypes:
-            rg = requires_grad and Adjoint.is_differentiable_value_type(dtype)
-            outputs.append(wp.empty(out_shape, dtype=dtype, requires_grad=rg, device=device))
-    elif len(out_dtypes) == 1 and is_array(out):
-        if not types_equal(out.dtype, out_dtypes[0]):
-            raise TypeError(
-                f"Output array dtype {type_repr(out.dtype)} does not match expected dtype {type_repr(out_dtypes[0])}"
-            )
-        if out.shape != out_shape:
-            raise TypeError(f"Output array shape {out.shape} does not match expected shape {out_shape}")
-        outputs = [out]
-    elif len(out_dtypes) > 1:
-        if isinstance(out, tuple) or isinstance(out, list):
-            if len(out) != len(out_dtypes):
-                raise TypeError(
-                    f"Number of provided output arrays ({len(out)}) does not match expected number of function outputs ({len(out_dtypes)})"
-                )
-            for i, a in enumerate(out):
-                if not types_equal(a.dtype, out_dtypes[i]):
-                    raise TypeError(
-                        f"Output array {i} dtype {type_repr(a.dtype)} does not match expected dtype {type_repr(out_dtypes[i])}"
-                    )
-                if a.shape != out_shape:
-                    raise TypeError(f"Output array {i} shape {a.shape} does not match expected shape {out_shape}")
-            outputs = list(out)
-        else:
-            raise TypeError(
-                f"Invalid output provided, expected {len(out_dtypes)} Warp arrays with shape {out_shape} and dtypes ({', '.join(type_repr(t) for t in out_dtypes)})"
-            )
-
-    # create code for a kernel
-    code = """def map_kernel({kernel_args}):
-    {tids} = wp.tid()
-    {load_args}
-    """
-    if len(outputs) == 1:
-        code += "__out_0[{tids}] = {func_name}({arg_names})"
-    else:
-        code += ", ".join(f"__o_{i}" for i in range(len(outputs)))
-        code += " = {func_name}({arg_names})\n"
-        for i in range(len(outputs)):
-            code += f"    __out_{i}" + "[{tids}]" + f" = __o_{i}\n"
-
-    tids = [f"__tid_{i}" for i in range(len(out_shape))]
-
-    load_args = []
-    kernel_args = []
-    for arg_name, input in zip(arg_names, inputs):
-        if is_array(input):
-            arr_name = f"{arg_name}_array"
-            array_type_name = type(input).__name__
-            kernel_args.append(
-                f"{arr_name}: wp.{array_type_name}(dtype={type_to_code(input.dtype)}, ndim={input.ndim})"
-            )
-            shape = input.shape
-            indices = []
-            for i in range(1, len(shape) + 1):
-                if shape[-i] == 1:
-                    indices.append("0")
-                else:
-                    indices.append(tids[-i])
-
-            load_args.append(f"{arg_name} = {arr_name}[{', '.join(reversed(indices))}]")
-        else:
-            kernel_args.append(f"{arg_name}: {type_to_code(type(input))}")
-    for i, o in enumerate(outputs):
-        array_type_name = type(o).__name__
-        kernel_args.append(f"__out_{i}: wp.{array_type_name}(dtype={type_to_code(o.dtype)}, ndim={o.ndim})")
-    code = code.format(
-        func_name=func_name,
-        kernel_args=", ".join(kernel_args),
-        arg_names=", ".join(arg_names),
-        tids=", ".join(tids),
-        load_args="\n    ".join(load_args),
-    )
-    namespace = {}
-    namespace.update({"wp": wp, "warp": wp, func_name: wp_func, "Any": Any})
-    namespace.update(referenced_modules)
-    exec(code, namespace)
-
-    kernel = wp.Kernel(namespace["map_kernel"], key="map_kernel", source=code, module=module)
-    if return_kernel:
-        return kernel
-
-    wp.launch(
-        kernel,
-        dim=out_shape,
-        inputs=inputs,
-        outputs=outputs,
-        block_dim=block_dim,
-        device=device,
-    )
-
-    if len(outputs) == 1:
-        o = outputs[0]
-    else:
-        o = outputs
-
-    return o
-
-
-# code snippet for invoking cProfile
-# cp = cProfile.Profile()
-# cp.enable()
-# for i in range(1000):
-#     self.state = self.integrator.forward(self.model, self.state, self.sim_dt)
-
-# cp.disable()
-# cp.print_stats(sort='tottime')
-# exit(0)
-
-
-# helper kernels for initializing NVDB volumes from a dense array
-@wp.kernel
-def copy_dense_volume_to_nano_vdb_v(volume: wp.uint64, values: wp.array(dtype=wp.vec3, ndim=3)):
-    i, j, k = wp.tid()
-    wp.volume_store_v(volume, i, j, k, values[i, j, k])
-
-
-@wp.kernel
-def copy_dense_volume_to_nano_vdb_f(volume: wp.uint64, values: wp.array(dtype=wp.float32, ndim=3)):
-    i, j, k = wp.tid()
-    wp.volume_store_f(volume, i, j, k, values[i, j, k])
-
-
-@wp.kernel
-def copy_dense_volume_to_nano_vdb_i(volume: wp.uint64, values: wp.array(dtype=wp.int32, ndim=3)):
-    i, j, k = wp.tid()
-    wp.volume_store_i(volume, i, j, k, values[i, j, k])
-
-
-# represent an edge between v0, v1 with connected faces f0, f1, and opposite vertex o0, and o1
-# winding is such that first tri can be reconstructed as {v0, v1, o0}, and second tri as { v1, v0, o1 }
-class MeshEdge:
-    def __init__(self, v0, v1, o0, o1, f0, f1):
-        self.v0 = v0  # vertex 0
-        self.v1 = v1  # vertex 1
-        self.o0 = o0  # opposite vertex 1
-        self.o1 = o1  # opposite vertex 2
-        self.f0 = f0  # index of tri1
-        self.f1 = f1  # index of tri2
-
-
-class MeshAdjacency:
-    def __init__(self, indices, num_tris):
-        # map edges (v0, v1) to faces (f0, f1)
-        self.edges = {}
-        self.indices = indices
-
-        for index, tri in enumerate(indices):
-            self.add_edge(tri[0], tri[1], tri[2], index)
-            self.add_edge(tri[1], tri[2], tri[0], index)
-            self.add_edge(tri[2], tri[0], tri[1], index)
-
-    def add_edge(self, i0, i1, o, f):  # index1, index2, index3, index of triangle
-        key = (min(i0, i1), max(i0, i1))
-        edge = None
-
-        if key in self.edges:
-            edge = self.edges[key]
-
-            if edge.f1 != -1:
-                print("Detected non-manifold edge")
-                return
-            else:
-                # update other side of the edge
-                edge.o1 = o
-                edge.f1 = f
-        else:
-            # create new edge with opposite yet to be filled
-            edge = MeshEdge(i0, i1, o, -1, f, -1)
-
-        self.edges[key] = edge
-
-
-def mem_report():  # pragma: no cover
-    def _mem_report(tensors, mem_type):
-        """Print the selected tensors of type
-        There are two major storage types in our major concern:
-            - GPU: tensors transferred to CUDA devices
-            - CPU: tensors remaining on the system memory (usually unimportant)
-        Args:
-            - tensors: the tensors of specified type
-            - mem_type: 'CPU' or 'GPU' in current implementation"""
-        total_numel = 0
-        total_mem = 0
-        visited_data = []
-        for tensor in tensors:
-            if tensor.is_sparse:
-                continue
-            # a data_ptr indicates a memory block allocated
-            data_ptr = tensor.storage().data_ptr()
-            if data_ptr in visited_data:
-                continue
-            visited_data.append(data_ptr)
-
-            numel = tensor.storage().size()
-            total_numel += numel
-            element_size = tensor.storage().element_size()
-            mem = numel * element_size / 1024 / 1024  # 32bit=4Byte, MByte
-            total_mem += mem
-        print(f"Type: {mem_type:<4} | Total Tensors: {total_numel:>8} | Used Memory: {total_mem:>8.2f} MB")
-
-    import gc
-
-    import torch
-
-    gc.collect()
-
-    LEN = 65
-    objects = gc.get_objects()
-    # print('%s\t%s\t\t\t%s' %('Element type', 'Size', 'Used MEM(MBytes)') )
-    tensors = [obj for obj in objects if torch.is_tensor(obj)]
-    cuda_tensors = [t for t in tensors if t.is_cuda]
-    host_tensors = [t for t in tensors if not t.is_cuda]
-    _mem_report(cuda_tensors, "GPU")
-    _mem_report(host_tensors, "CPU")
-    print("=" * LEN)
-
-
-class ScopedDevice:
-    """A context manager to temporarily change the current default device.
-
-    For CUDA devices, this context manager makes the device's CUDA context
-    current and restores the previous CUDA context on exit. This is handy when
-    running Warp scripts as part of a bigger pipeline because it avoids any side
-    effects of changing the CUDA context in the enclosed code.
-
-    Attributes:
-        device (Device): The device that will temporarily become the default
-          device within the context.
-        saved_device (Device): The previous default device. This is restored as
-          the default device on exiting the context.
-    """
-
-    def __init__(self, device: Devicelike):
-        """Initializes the context manager with a device.
-
-        Args:
-            device: The device that will temporarily become the default device
-              within the context.
-        """
-        self.device = wp.get_device(device)
-
-    def __enter__(self):
-        # save the previous default device
-        self.saved_device = self.device.runtime.default_device
-
-        # make this the default device
-        self.device.runtime.default_device = self.device
-
-        # make it the current CUDA device so that device alias "cuda" will evaluate to this device
-        self.device.context_guard.__enter__()
-
-        return self.device
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        # restore original CUDA context
-        self.device.context_guard.__exit__(exc_type, exc_value, traceback)
-
-        # restore original target device
-        self.device.runtime.default_device = self.saved_device
-
-
-class ScopedStream:
-    """A context manager to temporarily change the current stream on a device.
-
-    Attributes:
-        stream (Stream or None): The stream that will temporarily become the device's
-          default stream within the context.
-        saved_stream (Stream): The device's previous current stream. This is
-          restored as the device's current stream on exiting the context.
-        sync_enter (bool): Whether to synchronize this context's stream with
-          the device's previous current stream on entering the context.
-        sync_exit (bool): Whether to synchronize the device's previous current
-          with this context's stream on exiting the context.
-        device (Device): The device associated with the stream.
-    """
-
-    def __init__(self, stream: wp.Stream | None, sync_enter: bool = True, sync_exit: bool = False):
-        """Initializes the context manager with a stream and synchronization options.
-
-        Args:
-            stream: The stream that will temporarily become the device's
-              default stream within the context.
-            sync_enter (bool): Whether to synchronize this context's stream with
-              the device's previous current stream on entering the context.
-            sync_exit (bool): Whether to synchronize the device's previous current
-              with this context's stream on exiting the context.
-        """
-
-        self.stream = stream
-        self.sync_enter = sync_enter
-        self.sync_exit = sync_exit
-        if stream is not None:
-            self.device = stream.device
-            self.device_scope = ScopedDevice(self.device)
-
-    def __enter__(self):
-        if self.stream is not None:
-            self.device_scope.__enter__()
-            self.saved_stream = self.device.stream
-            self.device.set_stream(self.stream, self.sync_enter)
-
-        return self.stream
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        if self.stream is not None:
-            self.device.set_stream(self.saved_stream, self.sync_exit)
-            self.device_scope.__exit__(exc_type, exc_value, traceback)
-
-
-TIMING_KERNEL = 1
-TIMING_KERNEL_BUILTIN = 2
-TIMING_MEMCPY = 4
-TIMING_MEMSET = 8
-TIMING_GRAPH = 16
-TIMING_ALL = 0xFFFFFFFF
-
-
-# timer utils
-class ScopedTimer:
-    indent = -1
-
-    enabled = True
-
-    def __init__(
-        self,
-        name: str,
-        active: bool = True,
-        print: bool = True,
-        detailed: bool = False,
-        dict: dict[str, list[float]] | None = None,
-        use_nvtx: bool = False,
-        color: int | str = "rapids",
-        synchronize: bool = False,
-        cuda_filter: int = 0,
-        report_func: Callable[[list[TimingResult], str], None] | None = None,
-        skip_tape: bool = False,
-    ):
-        """Context manager object for a timer
-
-        Parameters:
-            name: Name of timer
-            active: Enables this timer
-            print: At context manager exit, print elapsed time to ``sys.stdout``
-            detailed: Collects additional profiling data using cProfile and calls ``print_stats()`` at context exit
-            dict: A dictionary of lists to which the elapsed time will be appended using ``name`` as a key
-            use_nvtx: If true, timing functionality is replaced by an NVTX range
-            color: ARGB value (e.g. 0x00FFFF) or color name (e.g. 'cyan') associated with the NVTX range
-            synchronize: Synchronize the CPU thread with any outstanding CUDA work to return accurate GPU timings
-            cuda_filter: Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
-            report_func: A callback function to print the activity report.
-              If ``None``,  :func:`wp.timing_print() <timing_print>` will be used.
-            skip_tape: If true, the timer will not be recorded in the tape
-
-        Attributes:
-            extra_msg (str): Can be set to a string that will be added to the printout at context exit.
-            elapsed (float): The duration of the ``with`` block used with this object
-            timing_results (list[TimingResult]): The list of activity timing results, if collection was requested using ``cuda_filter``
-        """
-        self.name = name
-        self.active = active and self.enabled
-        self.print = print
-        self.detailed = detailed
-        self.dict = dict
-        self.use_nvtx = use_nvtx
-        self.color = color
-        self.synchronize = synchronize
-        self.skip_tape = skip_tape
-        self.elapsed = 0.0
-        self.cuda_filter = cuda_filter
-        self.report_func = report_func or wp.timing_print
-        self.extra_msg = ""  # Can be used to add to the message printed at manager exit
-
-        if self.dict is not None:
-            if name not in self.dict:
-                self.dict[name] = []
-
-    def __enter__(self):
-        if not self.skip_tape and warp.context.runtime is not None and warp.context.runtime.tape is not None:
-            warp.context.runtime.tape.record_scope_begin(self.name)
-        if self.active:
-            if self.synchronize:
-                wp.synchronize()
-
-            if self.cuda_filter:
-                # begin CUDA activity collection, synchronizing if needed
-                timing_begin(self.cuda_filter, synchronize=not self.synchronize)
-
-            if self.detailed:
-                self.cp = cProfile.Profile()
-                self.cp.clear()
-                self.cp.enable()
-
-            if self.use_nvtx:
-                import nvtx
-
-                self.nvtx_range_id = nvtx.start_range(self.name, color=self.color)
-
-            if self.print:
-                ScopedTimer.indent += 1
-
-                if warp.config.verbose:
-                    indent = "    " * ScopedTimer.indent
-                    print(f"{indent}{self.name} ...", flush=True)
-
-            self.start = time.perf_counter_ns()
-
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        if not self.skip_tape and warp.context.runtime is not None and warp.context.runtime.tape is not None:
-            warp.context.runtime.tape.record_scope_end()
-        if self.active:
-            if self.synchronize:
-                wp.synchronize()
-
-            self.elapsed = (time.perf_counter_ns() - self.start) / 1000000.0
-
-            if self.use_nvtx:
-                import nvtx
-
-                nvtx.end_range(self.nvtx_range_id)
-
-            if self.detailed:
-                self.cp.disable()
-                self.cp.print_stats(sort="tottime")
-
-            if self.cuda_filter:
-                # end CUDA activity collection, synchronizing if needed
-                self.timing_results = timing_end(synchronize=not self.synchronize)
-            else:
-                self.timing_results = []
-
-            if self.dict is not None:
-                self.dict[self.name].append(self.elapsed)
-
-            if self.print:
-                indent = "    " * ScopedTimer.indent
-
-                if self.timing_results:
-                    self.report_func(self.timing_results, indent=indent)
-                    print()
-
-                if self.extra_msg:
-                    print(f"{indent}{self.name} took {self.elapsed:.2f} ms {self.extra_msg}")
-                else:
-                    print(f"{indent}{self.name} took {self.elapsed:.2f} ms")
-
-                ScopedTimer.indent -= 1
-
-
-# Allow temporarily enabling/disabling mempool allocators
-class ScopedMempool:
-    def __init__(self, device: Devicelike, enable: bool):
-        self.device = wp.get_device(device)
-        self.enable = enable
-
-    def __enter__(self):
-        self.saved_setting = wp.is_mempool_enabled(self.device)
-        wp.set_mempool_enabled(self.device, self.enable)
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        wp.set_mempool_enabled(self.device, self.saved_setting)
-
-
-# Allow temporarily enabling/disabling mempool access
-class ScopedMempoolAccess:
-    def __init__(self, target_device: Devicelike, peer_device: Devicelike, enable: bool):
-        self.target_device = target_device
-        self.peer_device = peer_device
-        self.enable = enable
-
-    def __enter__(self):
-        self.saved_setting = wp.is_mempool_access_enabled(self.target_device, self.peer_device)
-        wp.set_mempool_access_enabled(self.target_device, self.peer_device, self.enable)
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        wp.set_mempool_access_enabled(self.target_device, self.peer_device, self.saved_setting)
-
-
-# Allow temporarily enabling/disabling peer access
-class ScopedPeerAccess:
-    def __init__(self, target_device: Devicelike, peer_device: Devicelike, enable: bool):
-        self.target_device = target_device
-        self.peer_device = peer_device
-        self.enable = enable
-
-    def __enter__(self):
-        self.saved_setting = wp.is_peer_access_enabled(self.target_device, self.peer_device)
-        wp.set_peer_access_enabled(self.target_device, self.peer_device, self.enable)
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        wp.set_peer_access_enabled(self.target_device, self.peer_device, self.saved_setting)
-
-
-class ScopedCapture:
-    def __init__(self, device: Devicelike = None, stream=None, force_module_load=None, external=False):
-        self.device = device
-        self.stream = stream
-        self.force_module_load = force_module_load
-        self.external = external
-        self.active = False
-        self.graph = None
-
-    def __enter__(self):
-        try:
-            wp.capture_begin(
-                device=self.device, stream=self.stream, force_module_load=self.force_module_load, external=self.external
-            )
-            self.active = True
-            return self
-        except:
-            raise
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        if self.active:
-            try:
-                self.graph = wp.capture_end(device=self.device, stream=self.stream)
-            except Exception:
-                # Only report this exception if __exit__() was reached without an exception,
-                # otherwise re-raise the original exception.
-                if exc_type is None:
-                    raise
-            finally:
-                self.active = False
-
-
-def check_p2p():
-    """Check if the machine is configured properly for peer-to-peer transfers.
-
-    Returns:
-        A Boolean indicating whether the machine is configured properly for peer-to-peer transfers.
-        On Linux, this function attempts to determine if IOMMU is enabled and will return `False` if IOMMU is detected.
-        On other operating systems, it always return `True`.
-    """
-
-    # HACK: allow disabling P2P tests using an environment variable
-    disable_p2p_tests = os.getenv("WARP_DISABLE_P2P_TESTS", default="0")
-    if int(disable_p2p_tests):
-        return False
-
-    if sys.platform == "linux":
-        # IOMMU enablement can affect peer-to-peer transfers.
-        # On modern Linux, there should be IOMMU-related entries in the /sys file system.
-        # This should be more reliable than checking kernel logs like dmesg.
-        if os.path.isdir("/sys/class/iommu") and os.listdir("/sys/class/iommu"):
-            return False
-        if os.path.isdir("/sys/kernel/iommu_groups") and os.listdir("/sys/kernel/iommu_groups"):
-            return False
-
-    return True
-
-
-class timing_result_t(ctypes.Structure):
-    """CUDA timing struct for fetching values from C++"""
-
-    _fields_ = (
-        ("context", ctypes.c_void_p),
-        ("name", ctypes.c_char_p),
-        ("filter", ctypes.c_int),
-        ("elapsed", ctypes.c_float),
-    )
-
-
-class TimingResult:
-    """Timing result for a single activity."""
-
-    def __init__(self, device, name, filter, elapsed):
-        self.device: warp.context.Device = device
-        """The device where the activity was recorded."""
-
-        self.name: str = name
-        """The activity name."""
-
-        self.filter: int = filter
-        """The type of activity (e.g., ``warp.TIMING_KERNEL``)."""
-
-        self.elapsed: float = elapsed
-        """The elapsed time in milliseconds."""
-
-
-def timing_begin(cuda_filter: int = TIMING_ALL, synchronize: bool = True) -> None:
-    """Begin detailed activity timing.
-
-    Parameters:
-        cuda_filter: Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
-        synchronize: Whether to synchronize all CUDA devices before timing starts
-    """
-
-    if synchronize:
-        warp.synchronize()
-
-    warp.context.runtime.core.wp_cuda_timing_begin(cuda_filter)
-
-
-def timing_end(synchronize: bool = True) -> list[TimingResult]:
-    """End detailed activity timing.
-
-    Parameters:
-        synchronize: Whether to synchronize all CUDA devices before timing ends
-
-    Returns:
-        A list of :class:`TimingResult` objects for all recorded activities.
-    """
-
-    if synchronize:
-        warp.synchronize()
-
-    # get result count
-    count = warp.context.runtime.core.wp_cuda_timing_get_result_count()
-
-    # get result array from C++
-    result_buffer = (timing_result_t * count)()
-    warp.context.runtime.core.wp_cuda_timing_end(ctypes.byref(result_buffer), count)
-
-    # prepare Python result list
-    results = []
-    for r in result_buffer:
-        device = warp.context.runtime.context_map.get(r.context)
-        filter = r.filter
-        elapsed = r.elapsed
-
-        name = r.name.decode()
-        if filter == TIMING_KERNEL:
-            if name.endswith("forward"):
-                # strip trailing "_cuda_kernel_forward"
-                name = f"forward kernel {name[:-20]}"
-            else:
-                # strip trailing "_cuda_kernel_backward"
-                name = f"backward kernel {name[:-21]}"
-        elif filter == TIMING_KERNEL_BUILTIN:
-            if name.startswith("wp::"):
-                name = f"builtin kernel {name[4:]}"
-            else:
-                name = f"builtin kernel {name}"
-
-        results.append(TimingResult(device, name, filter, elapsed))
-
-    return results
-
-
-def timing_print(results: list[TimingResult], indent: str = "") -> None:
-    """Print timing results.
-
-    Parameters:
-        results: List of :class:`TimingResult` objects to print.
-        indent: Optional indentation to prepend to all output lines.
-    """
-
-    if not results:
-        print("No activity")
-        return
-
-    class Aggregate:
-        def __init__(self, count=0, elapsed=0):
-            self.count = count
-            self.elapsed = elapsed
-
-    device_totals = {}
-    activity_totals = {}
-
-    max_name_len = len("Activity")
-    for r in results:
-        name_len = len(r.name)
-        max_name_len = max(max_name_len, name_len)
-
-    activity_width = max_name_len + 1
-    activity_dashes = "-" * activity_width
-
-    print(f"{indent}CUDA timeline:")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    print(f"{indent}Time            | Device  | Activity")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    for r in results:
-        device_agg = device_totals.get(r.device.alias)
-        if device_agg is None:
-            device_totals[r.device.alias] = Aggregate(count=1, elapsed=r.elapsed)
-        else:
-            device_agg.count += 1
-            device_agg.elapsed += r.elapsed
-
-        activity_agg = activity_totals.get(r.name)
-        if activity_agg is None:
-            activity_totals[r.name] = Aggregate(count=1, elapsed=r.elapsed)
-        else:
-            activity_agg.count += 1
-            activity_agg.elapsed += r.elapsed
-
-        print(f"{indent}{r.elapsed:12.6f} ms | {r.device.alias:7s} | {r.name}")
-
-    print()
-    print(f"{indent}CUDA activity summary:")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    print(f"{indent}Total time      | Count   | Activity")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    for name, agg in activity_totals.items():
-        print(f"{indent}{agg.elapsed:12.6f} ms | {agg.count:7d} | {name}")
-
-    print()
-    print(f"{indent}CUDA device summary:")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    print(f"{indent}Total time      | Count   | Device")
-    print(f"{indent}----------------+---------+{activity_dashes}")
-    for device, agg in device_totals.items():
-        print(f"{indent}{agg.elapsed:12.6f} ms | {agg.count:7d} | {device}")
+    return get_deprecated_api(_utils, "wp", name)