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

warp/dlpack.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 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,450 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Python specification for DLpack:
-# https://dmlc.github.io/dlpack/latest/python_spec.html
+# TODO: Remove after cleaning up the public API.
 
-import ctypes
+from warp._src import dlpack as _dlpack
 
-import warp
-from warp.thirdparty.dlpack import (
-    DLDataType,
-    DLDataTypeCode,
-    DLDevice,
-    DLDeviceType,
-    DLManagedTensor,
-    _c_str_dltensor,
-)
 
-_c_str_used_dltensor = b"used_dltensor"
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
 
-PyMem_RawMalloc = ctypes.pythonapi.PyMem_RawMalloc
-PyMem_RawMalloc.argtypes = [ctypes.c_size_t]
-PyMem_RawMalloc.restype = ctypes.c_void_p
-
-PyMem_RawFree = ctypes.pythonapi.PyMem_RawFree
-PyMem_RawFree.argtypes = [ctypes.c_void_p]
-PyMem_RawFree.restype = None
-
-Py_IncRef = ctypes.pythonapi.Py_IncRef
-Py_IncRef.argtypes = [ctypes.py_object]
-Py_IncRef.restype = None
-
-Py_DecRef = ctypes.pythonapi.Py_DecRef
-Py_DecRef.argtypes = [ctypes.py_object]
-Py_DecRef.restype = None
-
-PyCapsule_Destructor = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
-
-PyCapsule_IsValid = ctypes.pythonapi.PyCapsule_IsValid
-PyCapsule_IsValid.argtypes = [ctypes.py_object, ctypes.c_char_p]
-PyCapsule_IsValid.restype = ctypes.c_int
-
-PyCapsule_GetPointer = ctypes.pythonapi.PyCapsule_GetPointer
-PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
-PyCapsule_GetPointer.restype = ctypes.c_void_p
-
-PyCapsule_SetName = ctypes.pythonapi.PyCapsule_SetName
-PyCapsule_SetName.argtypes = [ctypes.py_object, ctypes.c_char_p]
-PyCapsule_SetName.restype = ctypes.c_int
-
-
-class _DLPackTensorHolder:
-    """Class responsible for deleting DLManagedTensor memory after ownership is transferred from a capsule."""
-
-    def __new__(cls, *args, **kwargs):
-        instance = super().__new__(cls)
-        instance.mem_ptr = None
-        return instance
-
-    def __init__(self, mem_ptr):
-        self.mem_ptr = mem_ptr
-
-    def __del__(self):
-        if not self.mem_ptr:
-            return
-
-        managed_tensor = DLManagedTensor.from_address(self.mem_ptr)
-        if managed_tensor.deleter:
-            managed_tensor.deleter(self.mem_ptr)
-
-
-@ctypes.CFUNCTYPE(None, ctypes.c_void_p)
-def _dlpack_tensor_deleter(managed_ptr) -> None:
-    """A function to deallocate a DLManagedTensor."""
-
-    managed_tensor = DLManagedTensor.from_address(managed_ptr)
-
-    # unreference the source array
-    manager = ctypes.cast(managed_tensor.manager_ctx, ctypes.py_object)
-    ctypes.pythonapi.Py_DecRef(manager)
-
-    # free the DLManagedTensor memory, including shape and strides
-    PyMem_RawFree(ctypes.c_void_p(managed_ptr))
-
-
-@PyCapsule_Destructor
-def _dlpack_capsule_deleter(ptr) -> None:
-    """Destructor for a capsule holding a DLManagedTensor."""
-
-    capsule = ctypes.cast(ptr, ctypes.py_object)
-
-    if PyCapsule_IsValid(capsule, _c_str_dltensor):
-        managed_ptr = PyCapsule_GetPointer(capsule, _c_str_dltensor)
-        managed_tensor = DLManagedTensor.from_address(managed_ptr)
-        if managed_tensor.deleter:
-            managed_tensor.deleter(managed_ptr)
-
-
-def _device_to_dlpack(wp_device: warp.context.Device) -> DLDevice:
-    dl_device = DLDevice()
-
-    if wp_device.is_cpu:
-        dl_device.device_type = DLDeviceType.kDLCPU
-        dl_device.device_id = 0
-    elif wp_device.is_cuda:
-        dl_device.device_type = DLDeviceType.kDLCUDA
-        dl_device.device_id = wp_device.ordinal
-    else:
-        raise RuntimeError(f"Invalid device type converting to DLPack: {wp_device}")
-
-    return dl_device
-
-
-def device_to_dlpack(wp_device) -> DLDevice:
-    return _device_to_dlpack(warp.get_device(wp_device))
-
-
-def dtype_to_dlpack(wp_dtype) -> DLDataType:
-    if wp_dtype == warp.bool:
-        return (DLDataTypeCode.kDLBool, 8, 1)
-    if wp_dtype == warp.int8:
-        return (DLDataTypeCode.kDLInt, 8, 1)
-    elif wp_dtype == warp.uint8:
-        return (DLDataTypeCode.kDLUInt, 8, 1)
-    elif wp_dtype == warp.int16:
-        return (DLDataTypeCode.kDLInt, 16, 1)
-    elif wp_dtype == warp.uint16:
-        return (DLDataTypeCode.kDLUInt, 16, 1)
-    elif wp_dtype == warp.int32:
-        return (DLDataTypeCode.kDLInt, 32, 1)
-    elif wp_dtype == warp.uint32:
-        return (DLDataTypeCode.kDLUInt, 32, 1)
-    elif wp_dtype == warp.int64:
-        return (DLDataTypeCode.kDLInt, 64, 1)
-    elif wp_dtype == warp.uint64:
-        return (DLDataTypeCode.kDLUInt, 64, 1)
-    elif wp_dtype == warp.float16:
-        return (DLDataTypeCode.kDLFloat, 16, 1)
-    elif wp_dtype == warp.float32:
-        return (DLDataTypeCode.kDLFloat, 32, 1)
-    elif wp_dtype == warp.float64:
-        return (DLDataTypeCode.kDLFloat, 64, 1)
-    else:
-        raise RuntimeError(f"No conversion from Warp type {wp_dtype} to DLPack type")
-
-
-def dtype_from_dlpack(dl_dtype):
-    # unpack to tuple for easier comparison
-    dl_dtype = (dl_dtype.type_code.value, dl_dtype.bits)
-
-    if dl_dtype == (DLDataTypeCode.kDLUInt, 1):
-        raise RuntimeError("Warp does not support bit boolean types")
-    elif dl_dtype == (DLDataTypeCode.kDLInt, 8):
-        return warp.types.int8
-    elif dl_dtype == (DLDataTypeCode.kDLInt, 16):
-        return warp.types.int16
-    elif dl_dtype == (DLDataTypeCode.kDLInt, 32):
-        return warp.types.int32
-    elif dl_dtype == (DLDataTypeCode.kDLInt, 64):
-        return warp.types.int64
-    elif dl_dtype == (DLDataTypeCode.kDLUInt, 8):
-        return warp.types.uint8
-    elif dl_dtype == (DLDataTypeCode.kDLUInt, 16):
-        return warp.types.uint16
-    elif dl_dtype == (DLDataTypeCode.kDLUInt, 32):
-        return warp.types.uint32
-    elif dl_dtype == (DLDataTypeCode.kDLUInt, 64):
-        return warp.types.uint64
-    elif dl_dtype == (DLDataTypeCode.kDLFloat, 16):
-        return warp.types.float16
-    elif dl_dtype == (DLDataTypeCode.kDLFloat, 32):
-        return warp.types.float32
-    elif dl_dtype == (DLDataTypeCode.kDLFloat, 64):
-        return warp.types.float64
-    elif dl_dtype == (DLDataTypeCode.kDLComplex, 64):
-        raise RuntimeError("Warp does not support complex types")
-    elif dl_dtype == (DLDataTypeCode.kDLComplex, 128):
-        raise RuntimeError("Warp does not support complex types")
-    else:
-        raise RuntimeError(f"Unknown DLPack datatype {dl_dtype}")
-
-
-def device_from_dlpack(dl_device):
-    assert warp.context.runtime is not None, "Warp not initialized, call wp.init() before use"
-
-    if dl_device.device_type.value == DLDeviceType.kDLCPU or dl_device.device_type.value == DLDeviceType.kDLCUDAHost:
-        return warp.context.runtime.cpu_device
-    elif (
-        dl_device.device_type.value == DLDeviceType.kDLCUDA
-        or dl_device.device_type.value == DLDeviceType.kDLCUDAManaged
-    ):
-        return warp.context.runtime.cuda_devices[dl_device.device_id]
-    else:
-        raise RuntimeError(f"Unknown device type from DLPack: {dl_device.device_type.value}")
-
-
-def shape_to_dlpack(shape):
-    a = (ctypes.c_int64 * len(shape))(*shape)
-    return a
-
-
-def strides_to_dlpack(strides, dtype):
-    # convert from byte count to element count
-    ndim = len(strides)
-    a = (ctypes.c_int64 * ndim)()
-    dtype_size = warp.types.type_size_in_bytes(dtype)
-    for i in range(ndim):
-        a[i] = strides[i] // dtype_size
-    return a
-
-
-def to_dlpack(wp_array: warp.array):
-    """Convert a Warp array to another type of DLPack-compatible array.
-
-    Args:
-        wp_array: The source Warp array that will be converted.
-
-    Returns:
-        A capsule containing a DLManagedTensor that can be converted
-        to another array type without copying the underlying memory.
-    """
-
-    # DLPack does not support structured arrays
-    if isinstance(wp_array.dtype, warp.codegen.Struct):
-        raise RuntimeError("Cannot convert structured Warp arrays to DLPack.")
-
-    # handle vector types
-    if hasattr(wp_array.dtype, "_wp_scalar_type_"):
-        # vector type, flatten the dimensions into one tuple
-        target_dtype = wp_array.dtype._wp_scalar_type_
-        target_ndim = wp_array.ndim + len(wp_array.dtype._shape_)
-        target_shape = (*wp_array.shape, *wp_array.dtype._shape_)
-        dtype_strides = warp.types.strides_from_shape(wp_array.dtype._shape_, wp_array.dtype._wp_scalar_type_)
-        target_strides = (*wp_array.strides, *dtype_strides)
-    else:
-        # scalar type
-        target_dtype = wp_array.dtype
-        target_ndim = wp_array.ndim
-        target_shape = wp_array.shape
-        target_strides = wp_array.strides
-
-    if wp_array.pinned:
-        dl_device = DLDevice()
-        dl_device.device_type = DLDeviceType.kDLCUDAHost
-        dl_device.device_id = 0
-    else:
-        dl_device = _device_to_dlpack(wp_array.device)
-
-    # allocate DLManagedTensor, shape, and strides together
-    managed_tensor_size = ctypes.sizeof(DLManagedTensor)
-    padding = managed_tensor_size & 7
-    shape_size = target_ndim * 8
-    mem_size = managed_tensor_size + padding + 2 * shape_size
-    mem_ptr = PyMem_RawMalloc(mem_size)
-    assert mem_ptr, "Failed to allocate memory for DLManagedTensor"
-
-    # set managed tensor attributes
-    managed_tensor = DLManagedTensor.from_address(mem_ptr)
-    managed_tensor.dl_tensor.data = wp_array.ptr
-    managed_tensor.dl_tensor.device = dl_device
-    managed_tensor.dl_tensor.ndim = target_ndim
-    managed_tensor.dl_tensor.dtype = dtype_to_dlpack(target_dtype)
-    managed_tensor.dl_tensor.byte_offset = 0
-
-    # shape
-    shape_offset = managed_tensor_size + padding
-    shape_ptr = ctypes.cast(mem_ptr + shape_offset, ctypes.POINTER(ctypes.c_int64))
-    for i in range(target_ndim):
-        shape_ptr[i] = target_shape[i]
-    managed_tensor.dl_tensor.shape = shape_ptr
-
-    # strides, if not contiguous
-    if wp_array.is_contiguous:
-        managed_tensor.dl_tensor.strides = None
-    else:
-        stride_offset = shape_offset + shape_size
-        stride_ptr = ctypes.cast(mem_ptr + stride_offset, ctypes.POINTER(ctypes.c_int64))
-        dtype_size = warp.types.type_size_in_bytes(target_dtype)
-        for i in range(target_ndim):
-            stride_ptr[i] = target_strides[i] // dtype_size
-        managed_tensor.dl_tensor.strides = stride_ptr
-
-    # DLManagedTensor holds a reference to the source array
-    managed_tensor.manager_ctx = id(wp_array)
-    Py_IncRef(wp_array)
-
-    managed_tensor.deleter = _dlpack_tensor_deleter
-
-    # NOTE: jax.ffi.pycapsule() defines the PyCapsule_New() argtypes incorrectly, which causes problems.
-    # Here we make sure that the PyCapsule_Destructor callback is correctly defined.
-    PyCapsule_New = ctypes.pythonapi.PyCapsule_New
-    PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, PyCapsule_Destructor]
-    PyCapsule_New.restype = ctypes.py_object
-
-    capsule = PyCapsule_New(
-        ctypes.byref(managed_tensor),
-        _c_str_dltensor,
-        _dlpack_capsule_deleter,
-    )
-
-    return capsule
-
-
-def dtype_is_compatible(dl_dtype, wp_dtype):
-    if dl_dtype.bits % 8 != 0:
-        raise RuntimeError("Data types with less than 8 bits are not supported")
-
-    if dl_dtype.type_code.value == DLDataTypeCode.kDLFloat:
-        if dl_dtype.bits == 16:
-            return wp_dtype == warp.float16
-        elif dl_dtype.bits == 32:
-            return wp_dtype == warp.float32
-        elif dl_dtype.bits == 64:
-            return wp_dtype == warp.float64
-    elif dl_dtype.type_code.value == DLDataTypeCode.kDLInt or dl_dtype.type_code.value == DLDataTypeCode.kDLUInt:
-        if dl_dtype.bits == 8:
-            return wp_dtype == warp.int8 or wp_dtype == warp.uint8
-        elif dl_dtype.bits == 16:
-            return wp_dtype == warp.int16 or wp_dtype == warp.uint16
-        elif dl_dtype.bits == 32:
-            return wp_dtype == warp.int32 or wp_dtype == warp.uint32
-        elif dl_dtype.bits == 64:
-            return wp_dtype == warp.int64 or wp_dtype == warp.uint64
-    elif dl_dtype.type_code.value == DLDataTypeCode.kDLBfloat:
-        raise RuntimeError("Bfloat data type is not supported")
-    elif dl_dtype.type_code.value == DLDataTypeCode.kDLComplex:
-        raise RuntimeError("Complex data types are not supported")
-    else:
-        raise RuntimeError(f"Unsupported DLPack dtype {(str(dl_dtype.type_code), dl_dtype.bits)}")
-
-
-def _from_dlpack(capsule, dtype=None) -> warp.array:
-    """Convert a DLPack capsule into a Warp array without copying.
-
-    Args:
-        capsule: A DLPack capsule wrapping an external array or tensor.
-        dtype: An optional Warp data type to interpret the source data.
-
-    Returns:
-        A new Warp array that uses the same underlying memory as the input capsule.
-    """
-
-    assert PyCapsule_IsValid(capsule, _c_str_dltensor), "Invalid capsule"
-    mem_ptr = PyCapsule_GetPointer(capsule, _c_str_dltensor)
-    managed_tensor = DLManagedTensor.from_address(mem_ptr)
-
-    dlt = managed_tensor.dl_tensor
-
-    device = device_from_dlpack(dlt.device)
-    pinned = dlt.device.device_type.value == DLDeviceType.kDLCUDAHost
-    shape = tuple(dlt.shape[dim] for dim in range(dlt.ndim))
-
-    # strides, if not contiguous
-    itemsize = dlt.dtype.bits // 8
-    if dlt.strides:
-        strides = tuple(dlt.strides[dim] * itemsize for dim in range(dlt.ndim))
-    else:
-        strides = None
-
-    # handle multi-lane dtypes as another dimension
-    if dlt.dtype.lanes > 1:
-        shape = (*shape, dlt.dtype.lanes)
-        if strides is not None:
-            strides = (*strides, itemsize)
-
-    if dtype is None:
-        # automatically detect dtype
-        dtype = dtype_from_dlpack(dlt.dtype)
-
-    elif hasattr(dtype, "_wp_scalar_type_"):
-        # handle vector/matrix types
-
-        if not dtype_is_compatible(dlt.dtype, dtype._wp_scalar_type_):
-            raise RuntimeError(f"Incompatible data types: {dlt.dtype} and {dtype}")
-
-        dtype_shape = dtype._shape_
-        dtype_dims = len(dtype._shape_)
-        if dtype_dims > len(shape) or dtype_shape != shape[-dtype_dims:]:
-            raise RuntimeError(
-                f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, ensure that source inner shape is {dtype_shape}"
-            )
-
-        if strides is not None:
-            # ensure the inner strides are contiguous
-            stride = itemsize
-            for i in range(dtype_dims):
-                if strides[-i - 1] != stride:
-                    raise RuntimeError(
-                        f"Could not convert DLPack tensor with shape {shape} to Warp array with dtype={dtype}, because the source inner strides are not contiguous"
-                    )
-                stride *= dtype_shape[-i - 1]
-            strides = tuple(strides[:-dtype_dims]) or (itemsize,)
-
-        shape = tuple(shape[:-dtype_dims]) or (1,)
-
-    elif not dtype_is_compatible(dlt.dtype, dtype):
-        # incompatible dtype requested
-        raise RuntimeError(f"Incompatible data types: {dlt.dtype} and {dtype}")
-
-    a = warp.types.array(
-        ptr=dlt.data, dtype=dtype, shape=shape, strides=strides, copy=False, device=device, pinned=pinned
-    )
-
-    # take ownership of the DLManagedTensor
-    a._dlpack_tensor_holder = _DLPackTensorHolder(mem_ptr)
-
-    # rename the capsule so that it no longer owns the DLManagedTensor
-    PyCapsule_SetName(capsule, _c_str_used_dltensor)
-
-    return a
-
-
-def from_dlpack(source, dtype=None) -> warp.array:
-    """Convert a source array or DLPack capsule into a Warp array without copying.
-
-    Args:
-        source: A DLPack-compatible array or PyCapsule
-        dtype: An optional Warp data type to interpret the source data.
-
-    Returns:
-        A new Warp array that uses the same underlying memory as the input
-        pycapsule.
-    """
-
-    # See https://data-apis.org/array-api/2022.12/API_specification/generated/array_api.array.__dlpack__.html
-
-    if hasattr(source, "__dlpack__"):
-        device_type, device_id = source.__dlpack_device__()
-        # Check if the source lives on a CUDA device
-        if device_type in (DLDeviceType.kDLCUDA, DLDeviceType.kDLCUDAManaged):
-            # Assume that the caller will use the array on its device's current stream.
-            # Note that we pass 1 for the null stream, per DLPack spec.
-            cuda_stream = warp.get_cuda_device(device_id).stream.cuda_stream or 1
-        elif device_type == DLDeviceType.kDLCPU:
-            # No stream sync for CPU arrays.
-            cuda_stream = None
-        elif device_type == DLDeviceType.kDLCUDAHost:
-            # For pinned memory, we sync with the current CUDA device's stream.
-            # Note that we pass 1 for the null stream, per DLPack spec.
-            cuda_stream = warp.get_cuda_device().stream.cuda_stream or 1
-        else:
-            raise TypeError("Unsupported source device")
-
-        capsule = source.__dlpack__(stream=cuda_stream)
-
-    else:
-        # legacy behaviour, assume source is a capsule
-        capsule = source
-
-    return _from_dlpack(capsule, dtype=dtype)
+    return get_deprecated_api(_dlpack, "wp", name)

warp/examples/distributed/example_jacobi_mpi.py

@@ -39,7 +39,6 @@ import numpy as np
 from mpi4py import MPI
 
 import warp as wp
-import warp.context
 from warp.types import warp_type_to_np_dtype
 
 wp.config.quiet = True  # Suppress wp.init() output
@@ -50,7 +49,7 @@ wptype = wp.float32  # Global precision setting, can set wp.float64 here for dou
 pi = wptype(math.pi)  # GitHub #485
 
 
-def calc_default_device(mpi_comm: "MPI.Comm") -> warp.context.Device:
+def calc_default_device(mpi_comm: "MPI.Comm") -> wp.context.Device:
     """Return the device that should be used for the current rank.
 
     This function is used to ensure that multiple MPI ranks running on the same
@@ -72,7 +71,7 @@ def calc_default_device(mpi_comm: "MPI.Comm") -> warp.context.Device:
     local_size = local_mpi_comm.Get_size()
     local_rank = local_mpi_comm.Get_rank()
 
-    num_cuda_devices = warp.get_cuda_device_count()
+    num_cuda_devices = wp.get_cuda_device_count()
 
     if 1 < num_cuda_devices < local_size:
         raise RuntimeError(
@@ -81,9 +80,9 @@ def calc_default_device(mpi_comm: "MPI.Comm") -> warp.context.Device:
 
     if 1 < num_cuda_devices:
         # Get the device based on local_rank
-        return warp.get_cuda_device(local_rank)
+        return wp.get_cuda_device(local_rank)
     else:
-        return warp.get_device()
+        return wp.get_device()
 
 
 def calc_decomp_1d(total_points: int, rank: int, total_ranks: int) -> Tuple[int, int]:

warp/examples/fem/example_adaptive_grid.py

@@ -16,7 +16,7 @@
 ###########################################################################
 # Example Adaptive Grid
 #
-# Demonstrates using an adaptive grid to increase the simulation resolition
+# Demonstrates using an adaptive grid to increase the simulation resolution
 # near a collider boundary.
 #
 ###########################################################################

warp/examples/fem/example_apic_fluid.py

@@ -400,7 +400,7 @@ class Example:
         cell_volume = np.prod(cell_size)
 
         radius = np.max(cell_size) * 0.5
-        volume = np.prod(cell_volume) * packing_fraction
+        volume = cell_volume * packing_fraction
 
         rng = np.random.default_rng(42)
         points += 2.0 * radius * (rng.random(points.shape) - 0.5)

warp/examples/fem/example_burgers.py

@@ -172,7 +172,7 @@ class Example:
 
         if self.velocity_field.space.degree > 0:
             # Integration on cells (if not piecewise-constant)
-            fem.utils.array_axpy(
+            fem.linalg.array_axpy(
                 x=fem.integrate(
                     cell_transport_form,
                     fields={"u": trial_velocity, "v": self._test, "w": trial_velocity},
@@ -196,19 +196,19 @@ class Example:
 
         # tmp = v0 - dt * k1
         tmp = self.velocity_field.space.make_field()
-        fem.utils.array_axpy(y=tmp.dof_values, x=self.velocity_field.dof_values, alpha=1.0, beta=0.0)
-        fem.utils.array_axpy(y=tmp.dof_values, x=k1, alpha=-self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=self.velocity_field.dof_values, alpha=1.0, beta=0.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k1, alpha=-self.sim_dt, beta=1.0)
         k2 = self._velocity_delta(tmp)
 
         # tmp = v0 - dt * (0.25 * k1 + 0.25 * k2)
-        fem.utils.array_axpy(y=tmp.dof_values, x=k1, alpha=0.75 * self.sim_dt, beta=1.0)
-        fem.utils.array_axpy(y=tmp.dof_values, x=k2, alpha=-0.25 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k1, alpha=0.75 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=tmp.dof_values, x=k2, alpha=-0.25 * self.sim_dt, beta=1.0)
         k3 = self._velocity_delta(tmp)
 
         # v = v0 - dt * (1/6 * k1 + 1/6 * k2 + 2/3 * k3)
-        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k1, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
-        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k2, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
-        fem.utils.array_axpy(y=self.velocity_field.dof_values, x=k3, alpha=-2.0 / 3.0 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k1, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k2, alpha=-1.0 / 6.0 * self.sim_dt, beta=1.0)
+        fem.linalg.array_axpy(y=self.velocity_field.dof_values, x=k3, alpha=-2.0 / 3.0 * self.sim_dt, beta=1.0)
 
         # Apply slope limiter
         if self.velocity_field.space.degree > 0:

warp/examples/fem/example_diffusion.py

@@ -27,7 +27,7 @@
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-from warp.fem.utils import array_axpy
+from warp.fem.linalg import array_axpy
 
 
 @fem.integrand

warp/examples/fem/example_distortion_energy.py

@@ -183,7 +183,7 @@ class Example:
             fem_example_utils.bsr_cg(u_matrix, b=u_rhs, x=du, quiet=self._quiet)
 
             # Accumulate to UV field
-            fem.utils.array_axpy(x=du, y=self._u_field.dof_values, alpha=-1.0, beta=1.0)
+            fem.linalg.array_axpy(x=du, y=self._u_field.dof_values, alpha=-1.0, beta=1.0)
 
     def render(self):
         # Visualization

warp/examples/fem/example_mixed_elasticity.py

@@ -159,7 +159,7 @@ class Example:
             self._geo, degree=degree, dtype=wp.vec2, element_basis=fem.ElementBasis.SERENDIPITY
         )
 
-        if isinstance(self._geo.reference_cell(), fem.geometry.element.Triangle):
+        if self._geo.reference_cell() == fem.Element.TRIANGLE:
             # triangle elements
             tau_basis = fem.ElementBasis.NONCONFORMING_POLYNOMIAL
             tau_degree = degree - 1
@@ -240,7 +240,7 @@ class Example:
             # Extract result -- cast to float32 and accumulate to displacement field
             delta_u = wp.empty_like(self._u_field.dof_values)
             wp.utils.array_cast(in_array=x, out_array=delta_u)
-            fem.utils.array_axpy(x=delta_u, y=self._u_field.dof_values)
+            fem.linalg.array_axpy(x=delta_u, y=self._u_field.dof_values)
 
         # Evaluate area conservation, should converge to 1.0 as Poisson ratio approaches 1.0
         final_area = fem.integrate(

warp/examples/fem/example_navier_stokes.py

@@ -28,7 +28,7 @@
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-from warp.fem.utils import array_axpy
+from warp.fem.linalg import array_axpy
 
 
 @wp.func

warp/examples/fem/example_nonconforming_contact.py

@@ -123,7 +123,7 @@ class Example:
                 [
                     [1.0, poisson, 0.0],
                     [poisson, 1.0, 0.0],
-                    [0.0, 0.0, (2.0 * (1.0 + poisson)) * (1.0 - poisson * poisson)],
+                    [0.0, 0.0, (1.0 - poisson * poisson) / (2.0 * (1.0 + poisson))],
                 ]
             )
         )
@@ -142,14 +142,14 @@ class Example:
         # Store stress degrees of freedom as symmetric tensors (3 dof) rather than full 2x2 matrices
         self._tau1_space = fem.make_polynomial_space(
             self._geo1,
-            degree=degree - 1,
+            degree=degree,
             discontinuous=True,
             element_basis=fem.ElementBasis.LAGRANGE,
             dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
         )
         self._tau2_space = fem.make_polynomial_space(
             self._geo2,
-            degree=degree - 1,
+            degree=degree,
             discontinuous=True,
             element_basis=fem.ElementBasis.LAGRANGE,
             dof_mapper=fem.SymmetricTensorMapper(wp.mat22),
@@ -173,7 +173,7 @@ class Example:
 
         # Damped update of coupling stress (for stability)
         alpha = 0.1
-        fem.utils.array_axpy(
+        fem.linalg.array_axpy(
             x=self._sig2_field_new.dof_values, y=self._sig2_field.dof_values, alpha=alpha, beta=1.0 - alpha
         )
 
@@ -219,7 +219,7 @@ class Example:
         u_rhs = fem.integrate(gravity_form, fields={"v": u_test}, values={"gravity": gravity}, output_dtype=wp.vec2d)
 
         # Add boundary stress from other solid field
-        other_stress_field = fem.field.field.NonconformingField(boundary, other_stress_field)
+        other_stress_field = fem.NonconformingField(boundary, other_stress_field)
         fem.integrate(
             boundary_stress_form,
             fields={"u": u_bd_test, "tau": other_stress_field},
@@ -232,8 +232,8 @@ class Example:
             bottom_boundary_projector_form, fields={"u": u_bd_trial, "v": u_bd_test}, assembly="nodal"
         )
 
-        # read displacement from other body set create bottom boundary Dirichlet BC
-        other_u_field = fem.field.field.NonconformingField(boundary, other_u_field)
+        # displacement from other body defines bottom boundary Dirichlet BC
+        other_u_field = fem.NonconformingField(boundary, other_u_field)
         u_bd_rhs = fem.integrate(
             bottom_boundary_projector_form, fields={"u": other_u_field, "v": u_bd_test}, assembly="nodal"
         )

warp/examples/fem/example_stokes.py

@@ -27,7 +27,7 @@
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-from warp.fem.utils import array_axpy
+from warp.fem.linalg import array_axpy
 
 
 @fem.integrand

warp/examples/fem/example_stokes_transfer.py

@@ -29,7 +29,7 @@ import numpy as np
 import warp as wp
 import warp.examples.fem.utils as fem_example_utils
 import warp.fem as fem
-from warp.fem.utils import array_axpy
+from warp.fem.linalg import array_axpy
 from warp.utils import array_cast