diff --git a/.gitignore b/.gitignore
index a47b0458..129afbb7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -133,6 +133,7 @@ env.bak/
 venv.bak/
 .xlb-env/
 .xlb_install_test_venvs/
+.xlb_backend_test_venvs/
 
 # Spyder project settings
 .spyderproject
diff --git a/README.md b/README.md
index cbf56a0a..c2ad8eb7 100644
--- a/README.md
+++ b/README.md
@@ -9,15 +9,22 @@
 XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax), [NVIDIA Warp](https://github.com/NVIDIA/warp), and [Neon](https://github.com/Autodesk/Neon) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the Warp backend, XLB offers state-of-the-art single-GPU performance, and with the new Neon backend it extends to multi-GPU (single-resolution). More importantly, the Neon backend provides grid refinement capabilities for multi-resolution simulations.
 
 ## Getting Started
-To get started with XLB, you can install it using pip. There are different installation options depending on your hardware and needs:
+To get started with XLB, you can install it using pip. XLB imports Warp at import time and Warp must come from exactly one of two extras, so every install picks either `xlb[warp]` or `xlb[neon]`:
+
+| Extra | Source of Warp | Use for |
+| --- | --- | --- |
+| `xlb[warp]` | `warp-lang` from PyPI | JAX and Warp backends (single-GPU) |
+| `xlb[neon]` | the fork bundled in `neon_gpu` | Neon backend (multi-GPU / multi-resolution) |
+
+Both unpack into the same `site-packages/warp`, so installing both leaves whichever pip wrote last in charge; XLB warns on import if it finds both.
 
 ### Basic Installation (CPU-only)
 ```bash
-pip install xlb
+pip install "xlb[warp]"
 ```
 
 ### Installation with Warp support (single-GPU)
-For the NVIDIA Warp backend (single-GPU, state-of-the-art performance):
+For the NVIDIA Warp backend (single-GPU, state-of-the-art performance) the same extra applies:
 ```bash
 pip install "xlb[warp]"
 ```
@@ -25,13 +32,13 @@ pip install "xlb[warp]"
 ### Installation with CUDA support (for NVIDIA GPUs)
 This installation is for the JAX backend with CUDA support:
 ```bash
-pip install "xlb[cuda]"
+pip install "xlb[cuda,warp]"
 ```
 
 ### Installation with TPU support
 This installation is for the JAX backend with TPU support:
 ```bash
-pip install "xlb[tpu]"
+pip install "xlb[tpu,warp]"
 ```
 
 ### Installation with Neon support
@@ -47,7 +54,7 @@ pip install '.[neon]'
 
 **Requirements:** The Neon wheel supports **Python 3.11** to **Python 3.14** on **Linux x86_64** and **Linux ARM**. 
 
-**Note:** Neon uses a custom fork of warp.
+**Note:** Neon uses a custom fork of warp, which is why `xlb[neon]` does not install `warp-lang`. If the environment already has `warp-lang`, remove it first with `pip uninstall warp-lang`, then reinstall `.[neon]`.
 
 ### Notes:
 - For Mac users: Use the basic CPU installation command as JAX's GPU support is not available on MacOS
@@ -57,7 +64,15 @@ pip install '.[neon]'
 To install the latest development version from source:
 
 ```bash
-pip install git+https://github.com/Autodesk/XLB.git
+pip install "xlb[warp] @ git+https://github.com/Autodesk/XLB.git"
+```
+
+### Running the tests
+`tests/run_backend_test_envs.py` builds a virtualenv per backend, installs XLB from the repository root, and runs the suite in each:
+
+```bash
+python tests/run_backend_test_envs.py                 # warp and neon
+python tests/run_backend_test_envs.py --backend neon  # one backend
 ```
 
 The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md).
diff --git a/docs/index.md b/docs/index.md
index 5a74c0d1..5cfb836a 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -7,20 +7,20 @@
 # XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning
 
 🎉 **Exciting News!** 🎉 XLB version 0.2.0 has been released, featuring a complete rewrite of the library and introducing support for the NVIDIA Warp backend! 
-XLB can now be installed via pip: `pip install xlb`.
+XLB can now be installed via pip: `pip install "xlb[warp]"`.
 
 XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax) and [NVIDIA Warp](https://github.com/NVIDIA/warp) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the new Warp backend, XLB now offers state-of-the-art performance for even faster simulations.
 
 ## Getting Started
-To get started with XLB, you can install it using pip:
+To get started with XLB, you can install it using pip. Warp comes from either the `warp` extra (PyPI `warp-lang`) or the `neon` extra (the fork bundled with `neon_gpu`); pick exactly one:
 ```bash
-pip install xlb
+pip install "xlb[warp]"
 ```
 
 To install the latest development version from source:
 
 ```bash
-pip install git+https://github.com/Autodesk/XLB.git
+pip install "xlb[warp] @ git+https://github.com/Autodesk/XLB.git"
 ```
 
 The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md).
diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
index 1616d32a..d2acabe5 100644
--- a/examples/cfd/flow_past_sphere_3d.py
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -85,6 +85,7 @@ def bc_profile_jax():
         H_y = wp_dtype(grid_shape[1] - 1)  # Height in y direction
         H_z = wp_dtype(grid_shape[2] - 1)  # Height in z direction
         two = wp_dtype(2.0)
+        _u_normal_vec = wp.types.vector(length=1, dtype=wp_dtype)
 
         @wp.func
         def bc_profile_warp(index: wp.vec3i):
@@ -98,7 +99,7 @@ def bc_profile_warp(index: wp.vec3i):
             r_squared = (two * y_center / H_y) ** two + (two * z_center / H_z) ** two
 
             # Parabolic profile: u = u_max * (1 - r²)
-            return wp.vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1)
+            return _u_normal_vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared))
 
         return bc_profile_warp
 
diff --git a/examples/cfd/multires_flow_past_sphere_3d.py b/examples/cfd/multires_flow_past_sphere_3d.py
index 968e9cd4..fa1a396c 100644
--- a/examples/cfd/multires_flow_past_sphere_3d.py
+++ b/examples/cfd/multires_flow_past_sphere_3d.py
@@ -171,7 +171,8 @@ def bc_profile():
     one = dtype(1.0)
     zero = dtype(0.0)
     u_max_wp = dtype(u_max)
-    _u_vec = wp.vec(velocity_set.d, dtype=dtype)
+    _u_vec = wp.types.vector(length=velocity_set.d, dtype=dtype)
+    _u_normal_vec = wp.types.vector(length=1, dtype=dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
@@ -191,7 +192,7 @@ def bc_profile_warp(index: wp.vec3i):
         # return _u_vec(u_max_wp * wp.max(zero, one - r_squared), zero, zero)
 
         # For Regularized and ZouHe
-        return wp.vec(u_max_wp * wp.max(zero, one - r_squared), length=1)
+        return _u_normal_vec(u_max_wp * wp.max(zero, one - r_squared))
 
     return bc_profile_warp
 
diff --git a/examples/cfd/rotating_sphere_3d.py b/examples/cfd/rotating_sphere_3d.py
index 184408c3..d9be4385 100644
--- a/examples/cfd/rotating_sphere_3d.py
+++ b/examples/cfd/rotating_sphere_3d.py
@@ -115,7 +115,7 @@
 def bc_profile():
     """Build a Warp function returning the rotational wall velocity at a voxel."""
     dtype = precision_policy.compute_precision.wp_dtype
-    _u_vec = wp.vec(velocity_set.d, dtype=dtype)
+    _u_vec = wp.types.vector(length=velocity_set.d, dtype=dtype)
     angular_velocity = _u_vec(0.0, rot_rate, 0.0)
     origin_np = shift + diam / 2
     origin_wp = _u_vec(origin_np[0], origin_np[1], origin_np[2])
diff --git a/examples/ibm/airfoil_ibm.py b/examples/ibm/airfoil_ibm.py
index 62264a22..c3f2f2d1 100644
--- a/examples/ibm/airfoil_ibm.py
+++ b/examples/ibm/airfoil_ibm.py
@@ -67,10 +67,11 @@ def define_boundary_indices(grid, velocity_set):
 def bc_profile(precision_policy, grid_shape, u_max):
     dtype = precision_policy.store_precision.wp_dtype
     u_max_d = dtype(u_max)
+    _u_normal_vec = wp.types.vector(length=1, dtype=dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(dtype(u_max_d), length=1)
+        return _u_normal_vec(dtype(u_max_d))
 
     return bc_profile_warp
 
diff --git a/examples/ibm/sphere_ibm.py b/examples/ibm/sphere_ibm.py
index 834e5736..1679dfcc 100644
--- a/examples/ibm/sphere_ibm.py
+++ b/examples/ibm/sphere_ibm.py
@@ -54,10 +54,11 @@ def define_boundary_indices(grid, velocity_set):
 def bc_profile(precision_policy, grid_shape, u_max):
     _dtype = precision_policy.store_precision.wp_dtype
     u_max_d = _dtype(u_max)
+    _u_normal_vec = wp.types.vector(length=1, dtype=_dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(_dtype(u_max_d), length=1)
+        return _u_normal_vec(_dtype(u_max_d))
 
     return bc_profile_warp
 
diff --git a/examples/ibm/wind_turbine_ibm.py b/examples/ibm/wind_turbine_ibm.py
index c8c688e0..f23d8294 100644
--- a/examples/ibm/wind_turbine_ibm.py
+++ b/examples/ibm/wind_turbine_ibm.py
@@ -40,10 +40,11 @@ def define_boundary_indices(grid, velocity_set):
 def bc_profile(precision_policy, grid_shape, u_max):
     _dtype = precision_policy.store_precision.wp_dtype
     u_max_d = _dtype(u_max)
+    _u_normal_vec = wp.types.vector(length=1, dtype=_dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(u_max_d, length=1)
+        return _u_normal_vec(u_max_d)
 
     return bc_profile_warp
 
diff --git a/examples/ibm/windtunnel_ibm.py b/examples/ibm/windtunnel_ibm.py
index 8f62c9db..584abd52 100644
--- a/examples/ibm/windtunnel_ibm.py
+++ b/examples/ibm/windtunnel_ibm.py
@@ -230,6 +230,7 @@ def bc_profile(precision_policy, grid_shape, u_max):
     u_max_d = _dtype(u_max)
     L_y = _dtype(grid_shape[1] - 1)
     L_z = _dtype(grid_shape[2] - 1)
+    _u_normal_vec = wp.types.vector(length=1, dtype=_dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
@@ -239,7 +240,7 @@ def bc_profile_warp(index: wp.vec3i):
         z_center = z - (L_z / _dtype(2.0))
         r_sq = ((_dtype(2.0) * y_center) / L_y) ** _dtype(2.0) + ((_dtype(2.0) * z_center) / L_z) ** _dtype(2.0)
         velocity_x = u_max_d * wp.max(_dtype(0.0), _dtype(1.0) - r_sq)
-        return wp.vec(velocity_x, length=1)
+        return _u_normal_vec(velocity_x)
 
     return bc_profile_warp
 
@@ -559,11 +560,11 @@ def save_drag_coefficient(cd_values, filename):
 cy_list = [c[1] for c in wheel_centers]
 cz_list = [c[2] for c in wheel_centers]
 
-_wheel_starts = wp.constant(wp.vec(len(starts_list), dtype=int)(starts_list))
-_wheel_ends = wp.constant(wp.vec(len(ends_list), dtype=int)(ends_list))
-_wheel_centers_x = wp.constant(wp.vec(len(cx_list), dtype=float)(cx_list))
-_wheel_centers_y = wp.constant(wp.vec(len(cy_list), dtype=float)(cy_list))
-_wheel_centers_z = wp.constant(wp.vec(len(cz_list), dtype=float)(cz_list))
+_wheel_starts = wp.constant(wp.types.vector(length=len(starts_list), dtype=int)(starts_list))
+_wheel_ends = wp.constant(wp.types.vector(length=len(ends_list), dtype=int)(ends_list))
+_wheel_centers_x = wp.constant(wp.types.vector(length=len(cx_list), dtype=float)(cx_list))
+_wheel_centers_y = wp.constant(wp.types.vector(length=len(cy_list), dtype=float)(cy_list))
+_wheel_centers_z = wp.constant(wp.types.vector(length=len(cz_list), dtype=float)(cz_list))
 
 bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max)
 stepper = setup_stepper(grid, bc_list, omega)
diff --git a/examples/out_of_core/autodiff_lbm.py b/examples/out_of_core/autodiff_lbm.py
index a15412b2..4057444a 100644
--- a/examples/out_of_core/autodiff_lbm.py
+++ b/examples/out_of_core/autodiff_lbm.py
@@ -194,11 +194,13 @@ def backward(
     else:
         raise ValueError("Invalid velocity set")
 
-    # Initialize XLB
+    # Initialize XLB. This example differentiates through the solver with
+    # wp.Tape, so it needs the adjoint kernels that XLB otherwise skips.
     xlb.init(
         velocity_set=velocity_set,
         default_backend=compute_backend,
         default_precision_policy=precision_policy,
+        enable_backward=True,
     )
 
     # Make grid for constructing stepper
diff --git a/setup.py b/setup.py
index ed261843..af8ae46a 100644
--- a/setup.py
+++ b/setup.py
@@ -1,36 +1,21 @@
-import os
 import platform
-import subprocess
 import sys
 
 from setuptools import setup, find_packages
-from setuptools.command.install import install
 
-
-def _neon_extra_requested():
-    """Best-effort detection of [neon] extra from install invocation."""
-    for arg in sys.argv:
-        if "neon" in arg and ("[" in arg or "xlb" in arg):
-            return True
-    return False
-
-
-def _uninstall_warp_lang(*, reason: str) -> None:
-    """Uninstall the ``warp-lang`` distribution so Neon's bundled warp fork is used."""
-    if os.environ.get("XLB_NEON_SKIP_UNINSTALL_WARP", "").lower() in ("1", "true", "yes"):
-        return
-    print(f"[xlb] {reason}")
-    try:
-        subprocess.run(
-            [sys.executable, "-m", "pip", "uninstall", "warp-lang", "-y"],
-            check=False,
-            capture_output=True,
-        )
-    except Exception as exc:  # noqa: BLE001
-        print(f"[xlb] Warning: failed to uninstall warp-lang: {exc}")
-
-
-_NEON_VERSION = "0.5.2a1"
+# Warp is required at import time, but it is deliberately not in
+# ``install_requires``: it must come from exactly one of the two extras.
+# ``[warp]`` installs ``warp-lang`` from PyPI, while ``[neon]`` installs
+# ``neon_gpu``, which bundles its own Warp fork. Both unpack into
+# ``site-packages/warp``, so if pip installs both, whichever lands last owns the
+# directory and later uninstalling ``warp-lang`` deletes files Neon needs.
+#
+# Listing ``warp-lang`` as a core dependency made that collision unavoidable for
+# ``pip install xlb[neon]``. A setup.py hook cannot undo it either: pip installs
+# dependencies after building this project, and for wheel and PEP 660 editable
+# installs the ``install`` command never runs at all. Keeping it out of the core
+# dependencies is the only way ``[neon]`` reliably gets one Warp.
+_NEON_VERSION = "0.5.2a5"
 _NEON_RELEASE_URL = f"https://github.com/Autodesk/Neon/releases/download/v{_NEON_VERSION}"
 
 
@@ -47,33 +32,6 @@ def _neon_wheel_requirement():
     return req
 
 
-class InstallWithNeonHooks(install):
-    """Uninstall ``warp-lang`` before and after install when ``[neon]`` is requested.
-
-    * **Before** ``pip``/setuptools install dependencies: removes any previously
-      installed ``warp-lang`` so an older or PyPI build does not linger next to
-      Neon's fork (``neon_gpu`` ships its own warp).
-    * **After** install: removes the ``warp-lang`` pulled in by ``install_requires``,
-      leaving Neon's warp as the one on the path.
-
-    Only runs when installing from source (e.g. sdist or git). Wheel installs
-    do not run setup.py, so for ``pip install xlb[neon]`` from PyPI you may
-    need to run ``pip uninstall warp-lang`` first if it is already installed.
-    Set XLB_NEON_SKIP_UNINSTALL_WARP=1 to disable this behaviour.
-    """
-
-    def run(self):
-        if _neon_extra_requested():
-            _uninstall_warp_lang(
-                reason=("Removing any existing warp-lang before Neon install (neon_gpu provides its own warp fork)."),
-            )
-        install.run(self)
-        if _neon_extra_requested():
-            _uninstall_warp_lang(
-                reason=("Removing PyPI warp-lang after install (core deps); use the warp bundled with neon_gpu."),
-            )
-
-
 setup(
     name="xlb",
     version="0.3.2",
@@ -93,11 +51,13 @@ def run(self):
         "pydantic>=2.9.1",
         "ruff>=0.14.1",
         "jax>=0.8.2",  # Base JAX CPU-only requirement
-        "warp-lang>=1.10.0",  # Required at import time (core modules import warp)
         "nvtx>=0.2.0",  # NVTX ranges (e.g. nse_multires_stepper); listed in requirements.txt
     ],
     extras_require={
-        "warp": ["warp-lang>=1.10.0"],  # Kept for explicit `pip install xlb[warp]` / Neon uninstall hook docs
+        # Exactly one of [warp] and [neon] is required; see the note above.
+        # Newest stable warp-lang on PyPI, the closest release to the Warp fork
+        # neon_gpu bundles (1.19.0.dev0). Move to 1.19 once it is released.
+        "warp": ["warp-lang>=1.18,<1.19"],
         "cuda": ["jax[cuda13]>=0.8.2"],  # For CUDA installations (pip install -U "jax[cuda13]")
         "tpu": ["jax[tpu]>=0.8.2"],  # For TPU installations
         # h5py: MultiresIO / Neon multi-resolution export to HDF5 (see xlb.utils.mesher).
@@ -106,5 +66,4 @@ def run(self):
     },
     python_requires=">=3.11",
     dependency_links=["https://storage.googleapis.com/jax-releases/libtpu_releases.html"],
-    cmdclass={"install": InstallWithNeonHooks},
 )
diff --git a/tests/backends/test_neon_backend.py b/tests/backends/test_neon_backend.py
new file mode 100644
index 00000000..26691723
--- /dev/null
+++ b/tests/backends/test_neon_backend.py
@@ -0,0 +1,81 @@
+"""
+Smoke coverage for the Neon compute backend.
+
+The rest of the suite only exercises the JAX and WARP backends, and
+``tests/install/flow_past_sphere_3d_test.py`` skips Neon because the sphere's
+``HalfwayBounceBackBC`` has no Neon implementation. These tests cover the parts
+of the Neon path that need no boundary condition: ``neon.init()`` registering
+XLB's custom Warp types for NVRTC, grid and field allocation, the device-side
+fill and host read-back, and operator kernel construction.
+
+Skipped when ``neon`` is missing (an env built with the ``[warp]`` extra rather
+than ``[neon]``) or when no CUDA device is available.
+"""
+
+import pytest
+
+neon = pytest.importorskip("neon", reason="neon_gpu is not installed; install XLB with the [neon] extra")
+
+import warp as wp  # noqa: E402
+
+import xlb  # noqa: E402
+from xlb.compute_backend import ComputeBackend  # noqa: E402
+from xlb.grid import grid_factory  # noqa: E402
+from xlb.precision_policy import PrecisionPolicy  # noqa: E402
+
+GRID_SHAPE = (16, 16, 16)
+CARDINALITY = 19
+FILL_VALUE = 2.5
+
+
+@pytest.fixture(scope="module", autouse=True)
+def neon_env():
+    """Initialize XLB on the Neon backend once for the module.
+
+    ``xlb.init`` calls ``neon.init()``, which must run before any kernel using
+    Neon's custom types is built, so every test here shares one initialization.
+    """
+    wp.init()
+    if wp.get_cuda_device_count() == 0:
+        pytest.skip("the Neon backend requires a CUDA device")
+
+    precision_policy = PrecisionPolicy.FP32FP32
+    xlb.init(
+        velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=ComputeBackend.NEON),
+        default_backend=ComputeBackend.NEON,
+        default_precision_policy=precision_policy,
+    )
+
+
+def test_neon_grid_shape():
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    assert grid.shape == GRID_SHAPE, "Neon grid shape is incorrect"
+
+
+def test_neon_field_cardinality():
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    f = grid.create_field(cardinality=CARDINALITY)
+    assert f.cardinality == CARDINALITY, "Neon field cardinality is incorrect"
+
+
+def test_neon_field_fill_value():
+    """Fill on the device, then read back on the host.
+
+    Exercises allocation, the fill kernel, and the device-to-host transfer.
+    """
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    f = grid.create_field(cardinality=CARDINALITY, fill_value=FILL_VALUE)
+    f.update_host(0)
+
+    for idx in (neon.Index_3d(0, 0, 0), neon.Index_3d(1, 2, 3), neon.Index_3d(*(d - 1 for d in GRID_SHAPE))):
+        for card in (0, CARDINALITY // 2, CARDINALITY - 1):
+            assert f.read(idx, card) == pytest.approx(FILL_VALUE), "Neon field not initialized with fill_value"
+
+
+def test_neon_operator_kernel_construction():
+    """Operators must produce a Neon functional and container, not Warp ones."""
+    from xlb.operator.equilibrium import QuadraticEquilibrium
+
+    op = QuadraticEquilibrium(compute_backend=ComputeBackend.NEON)
+    assert op.neon_functional is not None, "Neon functional was not constructed"
+    assert op.neon_container is not None, "Neon container was not constructed"
diff --git a/tests/grids/test_grid_jax.py b/tests/grids/test_grid_jax.py
index 7255b7f6..fc1dd2a6 100644
--- a/tests/grids/test_grid_jax.py
+++ b/tests/grids/test_grid_jax.py
@@ -3,7 +3,7 @@
 import xlb
 from xlb.compute_backend import ComputeBackend
 from xlb.grid import grid_factory
-from jax.sharding import Mesh
+from jax.sharding import Mesh, PartitionSpec
 from jax.experimental import mesh_utils
 import jax.numpy as jnp
 
@@ -30,7 +30,7 @@ def test_jax_2d_grid_initialization(grid_size):
 
     assert f.shape == (9,) + grid_shape, "Field shape is incorrect"
     assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect"
-    assert f.sharding.spec == ("cardinality", "x", "y"), "PartitionSpec is incorrect"
+    assert f.sharding.spec == PartitionSpec("cardinality", "x", "y"), "PartitionSpec is incorrect"
 
 
 @pytest.mark.parametrize("grid_size", [50, 100, 150])
@@ -46,12 +46,7 @@ def test_jax_3d_grid_initialization(grid_size):
 
     assert f.shape == (9,) + grid_shape, "Field shape is incorrect"
     assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect"
-    assert f.sharding.spec == (
-        "cardinality",
-        "x",
-        "y",
-        "z",
-    ), "PartitionSpec is incorrect"
+    assert f.sharding.spec == PartitionSpec("cardinality", "x", "y", "z"), "PartitionSpec is incorrect"
 
 
 def test_jax_grid_create_field_fill_value():
diff --git a/tests/install/flow_past_sphere_3d_test.py b/tests/install/flow_past_sphere_3d_test.py
index 6ae302cc..8022b088 100644
--- a/tests/install/flow_past_sphere_3d_test.py
+++ b/tests/install/flow_past_sphere_3d_test.py
@@ -97,6 +97,7 @@ def bc_profile_jax():
         H_y_w = wp_dtype(grid_shape[1] - 1)
         H_z_w = wp_dtype(grid_shape[2] - 1)
         two = wp_dtype(2.0)
+        _u_normal_vec = wp.types.vector(length=1, dtype=wp_dtype)
 
         @wp.func
         def bc_profile_warp(index: wp.vec3i):
@@ -105,7 +106,7 @@ def bc_profile_warp(index: wp.vec3i):
             y_center = y - (H_y_w / two)
             z_center = z - (H_z_w / two)
             r_squared = (two * y_center / H_y_w) ** two + (two * z_center / H_z_w) ** two
-            return wp.vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1)
+            return _u_normal_vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared))
 
         return bc_profile_warp
 
diff --git a/tests/install/run_install_matrix.py b/tests/install/run_install_matrix.py
index 9a5241a3..9a5e85ca 100644
--- a/tests/install/run_install_matrix.py
+++ b/tests/install/run_install_matrix.py
@@ -119,29 +119,30 @@ def main() -> int:
     scenarios: list[tuple[str, str, str, str]] = [
         (
             "jax-cpu",
-            "JAX (CPU): editable install with test extra",
-            "[test]",
-            "Base JAX CPU jaxlib + test deps (pytest). Warp-lang is still pulled in as a core dependency.",
+            "JAX (CPU): editable install with warp + test extras",
+            "[warp,test]",
+            "Base JAX CPU jaxlib + test deps (pytest). The warp extra is required because XLB imports "
+            "warp at import time and warp-lang is no longer a core dependency.",
         ),
         (
             "jax-cuda",
-            "JAX[cuda]: editable install with cuda + test extras",
-            "[cuda,test]",
+            "JAX[cuda]: editable install with cuda + warp + test extras",
+            "[cuda,warp,test]",
             "Adds the cuda extra from setup (jax[cuda13] per setup.py). Requires a matching CUDA stack.",
         ),
         (
             "warp",
             "WARP: explicit [warp,test] extras",
             "[warp,test]",
-            "Explicit WARP extra plus test deps; core install already includes warp-lang.",
+            "WARP extra plus test deps; warp-lang comes from PyPI.",
         ),
         (
             "neon",
             "NEON: editable install with neon + test extras",
             "[neon,test]",
             "Installs neon_gpu wheel per setup.py (Linux x86_64/aarch64, Python 3.11+). "
-            "Includes h5py for Neon multires HDF5 export. "
-            "Uninstalls any existing warp-lang before install and PyPI warp-lang after (Neon's fork).",
+            "Includes h5py for Neon multires HDF5 export. Must not pull warp-lang: neon_gpu "
+            "provides the Warp fork, and both would own site-packages/warp.",
         ),
     ]
 
diff --git a/tests/install/test_neon_install_warp_cleanup.py b/tests/install/test_neon_install_warp_cleanup.py
index 927af050..15bf62f2 100644
--- a/tests/install/test_neon_install_warp_cleanup.py
+++ b/tests/install/test_neon_install_warp_cleanup.py
@@ -1,6 +1,11 @@
 """
-Verify that installing XLB with the ``[neon]`` extra uninstalls a pre-existing
-``warp-lang`` (Neon ships its own fork via ``neon_gpu``).
+Verify that ``pip install -e .[neon]`` leaves exactly one Warp on the path.
+
+``neon_gpu`` bundles a Warp fork at the same ``site-packages/warp`` as PyPI's
+``warp-lang``, so if both are installed whichever pip unpacked last owns the
+directory. XLB therefore keeps ``warp-lang`` out of ``install_requires`` and
+offers it only through the ``[warp]`` extra; these tests pin that down, since a
+mixed install still imports and passes the rest of the suite.
 
 Uses an isolated virtualenv and ``pip install -e``, same style as
 ``run_install_matrix.py``. Skips on platforms where Neon wheels are not built
@@ -9,7 +14,6 @@
 
 from __future__ import annotations
 
-import os
 import platform
 import subprocess
 import sys
@@ -19,6 +23,12 @@
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 
+requires_neon_wheel = pytest.mark.skipif(
+    sys.platform != "linux" or platform.machine() not in ("x86_64", "aarch64"),
+    reason="Neon wheels: Linux x86_64 / aarch64 only",
+)
+requires_py311 = pytest.mark.skipif(sys.version_info < (3, 11), reason="XLB requires Python >= 3.11")
+
 
 def _venv_python(venv_dir: Path) -> Path:
     if sys.platform == "win32":
@@ -26,63 +36,78 @@ def _venv_python(venv_dir: Path) -> Path:
     return venv_dir / "bin" / "python"
 
 
-def _neon_wheels_supported() -> bool:
-    if sys.platform != "linux":
-        return False
-    return platform.machine() in ("x86_64", "aarch64")
+def _run(cmd: list[str], *, cwd: Path = REPO_ROOT) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=False)
 
 
-def _run(cmd: list[str], *, cwd: Path, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
-    return subprocess.run(
-        cmd,
-        cwd=cwd,
-        env=env,
-        capture_output=True,
-        text=True,
-        check=False,
-    )
-
-
-@pytest.mark.skipif(sys.version_info < (3, 11), reason="XLB requires Python >= 3.11")
-@pytest.mark.skipif(not _neon_wheels_supported(), reason="Neon wheels: Linux x86_64 / aarch64 only")
-def test_neon_editable_install_after_prior_warp_lang(tmp_path: Path) -> None:
-    """Pre-install ``warp-lang``, then ``pip install -e .[neon,test]``; imports must work."""
-    venv_dir = tmp_path / "venv"
+def _make_venv(venv_dir: Path) -> Path:
     subprocess.run([sys.executable, "-m", "venv", str(venv_dir)], check=True)
     venv_py = _venv_python(venv_dir)
     assert venv_py.is_file(), f"missing venv python: {venv_py}"
+    proc = _run([str(venv_py), "-m", "pip", "install", "--upgrade", "pip", "wheel"])
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    return venv_py
 
-    env = {**os.environ, "XLB_NEON_SKIP_UNINSTALL_WARP": ""}
 
-    proc = _run([str(venv_py), "-m", "pip", "install", "--upgrade", "pip", "wheel"], cwd=REPO_ROOT, env=env)
+def _pip_install(venv_py: Path, *args: str) -> None:
+    proc = _run([str(venv_py), "-m", "pip", "install", *args])
     assert proc.returncode == 0, proc.stdout + proc.stderr
 
-    proc = _run(
-        [str(venv_py), "-m", "pip", "install", "warp-lang==1.10.0"],
-        cwd=REPO_ROOT,
-        env=env,
-    )
-    assert proc.returncode == 0, proc.stdout + proc.stderr
 
-    proc = _run(
-        [str(venv_py), "-m", "pip", "install", "-e", ".[neon,test]"],
-        cwd=REPO_ROOT,
-        env=env,
+def _distribution_version(venv_py: Path, distribution: str) -> str | None:
+    """Return the installed version of *distribution*, or None if absent."""
+    code = (
+        "from importlib.metadata import PackageNotFoundError, version\n"
+        "try:\n"
+        f"    print(version({distribution!r}))\n"
+        "except PackageNotFoundError:\n"
+        "    print('')\n"
     )
+    proc = _run([str(venv_py), "-c", code])
     assert proc.returncode == 0, proc.stdout + proc.stderr
+    return proc.stdout.strip() or None
+
+
+@requires_py311
+@requires_neon_wheel
+def test_neon_install_does_not_pull_warp_lang(tmp_path: Path) -> None:
+    """A clean ``[neon]`` install must get its Warp from Neon and nowhere else."""
+    venv_py = _make_venv(tmp_path / "venv")
+    _pip_install(venv_py, "-e", ".[neon,test]")
 
-    proc = _run(
-        [str(venv_py), "-c", "import neon; import warp; print('ok', neon.__file__, warp.__file__)"],
-        cwd=REPO_ROOT,
-        env=env,
+    assert _distribution_version(venv_py, "neon_gpu") is not None, "neon_gpu was not installed"
+    warp_lang = _distribution_version(venv_py, "warp-lang")
+    assert warp_lang is None, (
+        f"warp-lang {warp_lang} came in alongside neon_gpu; both own site-packages/warp, "
+        "so whichever pip unpacked last wins. Keep warp-lang out of install_requires."
     )
-    assert proc.returncode == 0, proc.stdout + proc.stderr
-    assert "ok" in proc.stdout
 
-    # Repeat editable install: pre/post uninstall hooks must stay safe on re-run.
-    proc = _run(
-        [str(venv_py), "-m", "pip", "install", "-e", ".[neon,test]"],
-        cwd=REPO_ROOT,
-        env=env,
+    proc = _run([str(venv_py), "-c", "import neon, warp; print(warp.__version__, warp.__file__)"])
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "dev" in proc.stdout, f"expected Neon's Warp fork, got: {proc.stdout.strip()}"
+
+    # A repeated editable install must not disturb the result.
+    _pip_install(venv_py, "-e", ".[neon,test]")
+    assert _distribution_version(venv_py, "warp-lang") is None, "warp-lang reappeared on reinstall"
+
+
+@requires_py311
+@requires_neon_wheel
+def test_preexisting_warp_lang_is_reported(tmp_path: Path) -> None:
+    """pip will not remove a warp-lang that predates the install, so XLB must warn."""
+    venv_py = _make_venv(tmp_path / "venv")
+    _pip_install(venv_py, "warp-lang==1.10.0")
+    _pip_install(venv_py, "-e", ".[neon,test]")
+
+    code = (
+        "import warnings\n"
+        "with warnings.catch_warnings(record=True) as caught:\n"
+        "    warnings.simplefilter('always')\n"
+        "    import xlb\n"
+        "print([str(w.message) for w in caught])\n"
     )
+    proc = _run([str(venv_py), "-c", code])
     assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "warp-lang" in proc.stdout and "neon_gpu" in proc.stdout, (
+        f"importing XLB with both warp-lang and neon_gpu should warn, got: {proc.stdout.strip()}"
+    )
diff --git a/tests/run_backend_test_envs.py b/tests/run_backend_test_envs.py
new file mode 100755
index 00000000..f252c047
--- /dev/null
+++ b/tests/run_backend_test_envs.py
@@ -0,0 +1,341 @@
+#!/usr/bin/env python3
+"""
+Build a fresh virtualenv per compute backend, install XLB from the repository
+root, and run the test suite in each one.
+
+Usage (from the repository root)::
+
+    python tests/run_backend_test_envs.py                 # warp and neon
+    python tests/run_backend_test_envs.py --backend neon  # just one
+    python tests/run_backend_test_envs.py --reuse         # keep existing venvs
+    python tests/run_backend_test_envs.py --jax cuda      # GPU jaxlib in both envs
+
+Profiles
+--------
+warp
+    ``pip install -e .[warp,test]``. Warp comes from PyPI (``warp-lang``). The
+    suite's WARP and JAX tests run; the Neon tests in ``tests/backends`` skip
+    themselves because ``neon`` is not importable.
+neon
+    ``pip install -e .[neon,test]``. Installs the ``neon_gpu`` wheel, which
+    bundles its own Warp fork, and XLB's setup hook removes ``warp-lang`` so
+    that fork is the only Warp on the path. The same WARP and JAX tests
+    therefore run *against Neon's bundled Warp*, and the Neon tests also run.
+
+JAX is a core dependency (``install_requires``), so both profiles get it
+regardless of the backend; ``--jax cuda`` adds setup.py's ``cuda`` extra to
+swap the CPU-only jaxlib for a GPU one.
+
+Options
+-------
+--backend {warp,neon,both}  Profiles to build and test (default: both).
+--jax {cpu,cuda}            JAX flavor to install in every env (default: cpu).
+--venv-root PATH            Where to create the venvs
+                            (default: .xlb_backend_test_venvs, gitignored).
+--python PATH               Interpreter used to create the venvs. The [neon]
+                            extra resolves a wheel matching its version, so use
+                            this to test another CPython (default: this one).
+--reuse                     Reuse existing venvs instead of recreating them.
+--skip-install              Only run tests, assuming the venvs are populated.
+--skip-install-tests        Skip tests/install, whose Neon cleanup test builds
+                            another venv and downloads the neon_gpu wheel.
+--tests PATH [PATH ...]     Test paths to run (default: tests).
+--pytest-args "ARGS"        Extra pytest flags. Paths belong in --tests; these
+                            are appended, so a path here widens the selection.
+
+Environment
+-----------
+XLB_BACKEND_VENV_ROOT       Same as --venv-root.
+
+Exits non-zero if any install or test run fails.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shlex
+import shutil
+import subprocess
+import sys
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+# pytest reports "no tests collected" rather than success when every selected
+# test is skipped, which is what --tests tests/backends does in the warp env,
+# where the Neon tests skip themselves. Not a failure.
+PYTEST_NO_TESTS_COLLECTED = 5
+
+PROFILES: dict[str, dict[str, str]] = {
+    "warp": {
+        "extras": "warp,test",
+        "summary": "Warp backend, warp-lang from PyPI",
+    },
+    "neon": {
+        "extras": "neon,test",
+        "summary": "Neon backend, neon_gpu wheel with its bundled Warp fork",
+    },
+}
+
+
+@dataclass
+class ProfileResult:
+    name: str
+    extras: str
+    install_ok: bool = False
+    install_output: str = ""
+    tests_ok: bool | None = None
+    tests_exit: int | None = None
+    stack: dict[str, str] = field(default_factory=dict)
+    warning: str = ""
+    seconds: float = 0.0
+
+
+def venv_python(venv_dir: Path) -> Path:
+    if sys.platform == "win32":
+        return venv_dir / "Scripts" / "python.exe"
+    return venv_dir / "bin" / "python"
+
+
+def create_venv(venv_dir: Path, base_python: str, *, reuse: bool) -> Path:
+    py = venv_python(venv_dir)
+    if reuse and py.is_file():
+        print(f"venv: reusing {venv_dir}")
+        return py
+    if venv_dir.exists():
+        shutil.rmtree(venv_dir)
+    venv_dir.parent.mkdir(parents=True, exist_ok=True)
+    print(f"venv: creating {venv_dir} with {base_python}")
+    subprocess.run([base_python, "-m", "venv", str(venv_dir)], check=True)
+    return py
+
+
+def install_xlb(py: Path, repo_root: Path, extras: str) -> tuple[bool, str]:
+    """Editable-install XLB from *repo_root* with *extras*, capturing output."""
+    steps = [
+        [str(py), "-m", "pip", "install", "--upgrade", "pip", "wheel"],
+        [str(py), "-m", "pip", "install", "-e", f".[{extras}]"],
+    ]
+    transcript = ""
+    for cmd in steps:
+        print(f"pip: {' '.join(cmd[3:])}")
+        proc = subprocess.run(cmd, cwd=repo_root, capture_output=True, text=True, check=False)
+        transcript += proc.stdout + proc.stderr
+        if proc.returncode != 0:
+            return False, transcript
+    return True, transcript
+
+
+def describe_stack(py: Path, repo_root: Path) -> dict[str, str]:
+    """Report which Warp, Neon and JAX the env actually resolved to.
+
+    A wheel's version string alone does not identify Neon's Warp fork, so the
+    install location is reported too.
+    """
+    probe = """
+import importlib.util, json, os
+out = {}
+
+def where(mod):
+    spec = importlib.util.find_spec(mod)
+    return os.path.dirname(spec.origin) if spec and spec.origin else ""
+
+try:
+    import warp
+    out["warp"] = warp.__version__
+    out["warp_path"] = where("warp")
+except Exception as exc:
+    out["warp"] = f"unavailable ({exc})"
+
+out["neon"] = "importable" if importlib.util.find_spec("neon") else "absent"
+out["warp_lang_dist"] = "absent"
+try:
+    from importlib.metadata import version
+    out["warp_lang_dist"] = version("warp-lang")
+except Exception:
+    pass
+
+try:
+    import jax
+    out["jax"] = jax.__version__
+    out["jax_devices"] = ", ".join(sorted({d.platform for d in jax.devices()}))
+except Exception as exc:
+    out["jax"] = f"unavailable ({exc})"
+
+print(json.dumps(out))
+"""
+    proc = subprocess.run([str(py), "-c", probe], cwd=repo_root, capture_output=True, text=True, check=False)
+    if proc.returncode != 0:
+        return {"error": (proc.stderr or proc.stdout).strip()[-400:]}
+    import json
+
+    try:
+        return json.loads(proc.stdout.strip().splitlines()[-1])
+    except Exception:
+        return {"error": "could not parse environment probe output"}
+
+
+def warn_on_mixed_warp(stack: dict[str, str]) -> str:
+    """Flag a Neon env that also has the ``warp-lang`` distribution registered.
+
+    ``neon_gpu`` ships its own Warp fork at the same ``site-packages/warp``
+    path. XLB's ``[neon]`` extra no longer pulls ``warp-lang`` in, but pip does
+    not remove one that was installed beforehand, so a reused venv can still
+    hold both. Whichever pip unpacked last owns the files, and a later
+    ``pip uninstall warp-lang`` would delete files Neon needs. Tests pass either
+    way, so this is only detectable by looking.
+    """
+    if stack.get("neon") != "importable" or stack.get("warp_lang_dist") == "absent":
+        return ""
+    message = (
+        f"WARNING: both neon_gpu and warp-lang {stack['warp_lang_dist']} are installed.\n"
+        f"         The imported warp is {stack.get('warp', '?')}, so Neon's fork currently owns\n"
+        f"         site-packages/warp, but the env is order-dependent: uninstalling or\n"
+        f"         upgrading warp-lang would overwrite or delete Neon's Warp."
+    )
+    print(message)
+    return message
+
+
+def run_pytest(py: Path, repo_root: Path, *, targets: list[str], skip_install_tests: bool, extra_args: list[str]) -> int:
+    """Run the suite with output streamed, so progress is visible."""
+    cmd = [str(py), "-m", "pytest", *targets, "-q"]
+    if skip_install_tests:
+        cmd += ["--ignore", "tests/install"]
+    cmd += extra_args
+    print(f"pytest: {' '.join(cmd[1:])}\n")
+    return subprocess.run(cmd, cwd=repo_root, check=False).returncode
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Create a venv per backend, install XLB from the repo root, and run the tests.",
+    )
+    parser.add_argument("--backend", choices=[*PROFILES, "both"], default="both")
+    parser.add_argument("--jax", choices=["cpu", "cuda"], default="cpu", help="JAX flavor for every env")
+    parser.add_argument(
+        "--venv-root",
+        type=Path,
+        default=None,
+        help="Directory holding the per-backend venvs",
+    )
+    parser.add_argument("--python", default=sys.executable, help="Interpreter used to create the venvs")
+    parser.add_argument("--repo-root", type=Path, default=Path(__file__).resolve().parents[1])
+    parser.add_argument("--reuse", action="store_true", help="Reuse existing venvs")
+    parser.add_argument("--skip-install", action="store_true", help="Run tests only")
+    parser.add_argument("--skip-install-tests", action="store_true", help="Skip tests/install")
+    parser.add_argument(
+        "--tests",
+        nargs="+",
+        default=["tests"],
+        metavar="PATH",
+        help="Test paths to run (default: tests)",
+    )
+    parser.add_argument("--pytest-args", default="", help="Extra pytest flags; use --tests to narrow paths")
+    args = parser.parse_args()
+
+    # Subprocesses inherit this stdout and write straight to it, so keep our own
+    # prints line-buffered or they would surface out of order once redirected.
+    sys.stdout.reconfigure(line_buffering=True)
+
+    repo_root: Path = args.repo_root.resolve()
+    if not (repo_root / "setup.py").is_file():
+        print(f"ERROR: no setup.py in {repo_root}; pass --repo-root", file=sys.stderr)
+        return 2
+
+    default_root = os.environ.get("XLB_BACKEND_VENV_ROOT", repo_root / ".xlb_backend_test_venvs")
+    venv_root = Path(args.venv_root or default_root).resolve()
+    selected = list(PROFILES) if args.backend == "both" else [args.backend]
+    extra_pytest = shlex.split(args.pytest_args)
+
+    results: list[ProfileResult] = []
+    for name in selected:
+        extras = PROFILES[name]["extras"]
+        if args.jax == "cuda":
+            extras += ",cuda"
+        result = ProfileResult(name=name, extras=extras)
+        started = time.time()
+
+        print(f"\n{'=' * 78}")
+        print(f"Profile: {name} — {PROFILES[name]['summary']}")
+        print(f"Install: pip install -e .[{extras}]")
+        print(f"{'=' * 78}")
+
+        try:
+            py = create_venv(venv_root / name, args.python, reuse=args.reuse or args.skip_install)
+        except (subprocess.CalledProcessError, OSError) as exc:
+            result.install_output = f"venv creation failed: {exc}"
+            print(result.install_output, file=sys.stderr)
+            results.append(result)
+            continue
+
+        if args.skip_install:
+            result.install_ok = True
+            print("pip: skipped (--skip-install)")
+        else:
+            result.install_ok, result.install_output = install_xlb(py, repo_root, extras)
+            if not result.install_ok:
+                print("PIP INSTALL FAILED\n" + result.install_output[-4000:], file=sys.stderr)
+                result.seconds = time.time() - started
+                results.append(result)
+                continue
+            print("pip: ok")
+
+        result.stack = describe_stack(py, repo_root)
+        for key, value in result.stack.items():
+            print(f"  {key}: {value}")
+        result.warning = warn_on_mixed_warp(result.stack)
+        print()
+
+        result.tests_exit = run_pytest(
+            py,
+            repo_root,
+            targets=args.tests,
+            skip_install_tests=args.skip_install_tests,
+            extra_args=extra_pytest,
+        )
+        if result.tests_exit == PYTEST_NO_TESTS_COLLECTED:
+            result.tests_ok = None
+            verdict = "no tests ran (all skipped for this profile)"
+        else:
+            result.tests_ok = result.tests_exit == 0
+            verdict = "OK" if result.tests_ok else "FAIL"
+        result.seconds = time.time() - started
+        print(f"\npytest exit code: {result.tests_exit} ({verdict})")
+        results.append(result)
+
+    print(f"\n{'=' * 78}")
+    print("SUMMARY — XLB per-backend environments")
+    print(f"{'=' * 78}")
+    print(f"Repository: {repo_root}")
+    print(f"Venv root:  {venv_root}\n")
+
+    width = max((len(r.name) for r in results), default=8)
+    print(f"{'profile':<{width}}  {'install':^8}  {'tests':^6}  {'minutes':^8}  stack")
+    print("-" * 78)
+    for r in results:
+        install_s = "ok" if r.install_ok else "FAIL"
+        tests_s = "—" if r.tests_ok is None else ("ok" if r.tests_ok else "FAIL")
+        warp_desc = r.stack.get("warp", "?")
+        if r.stack.get("neon") == "importable":
+            warp_desc += " (neon-bundled)" if r.stack.get("warp_lang_dist") == "absent" else " (+warp-lang!)"
+        print(f"{r.name:<{width}}  {install_s:^8}  {tests_s:^6}  {r.seconds / 60:^8.1f}  warp {warp_desc}, jax {r.stack.get('jax', '?')}")
+
+    warned = [r for r in results if r.warning]
+    if warned:
+        print()
+        for r in warned:
+            print(f"{r.name}:\n{r.warning}")
+
+    failed = [r for r in results if not r.install_ok or r.tests_ok is False]
+    print()
+    if failed:
+        print(f"Overall: FAILURE ({', '.join(r.name for r in failed)})")
+        return 1
+    print("Overall: SUCCESS")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/xlb/__init__.py b/xlb/__init__.py
index 02b7a994..ce81bfde 100644
--- a/xlb/__init__.py
+++ b/xlb/__init__.py
@@ -1,3 +1,4 @@
+import warnings
 from importlib.metadata import PackageNotFoundError, version
 
 try:
@@ -5,6 +6,48 @@
 except PackageNotFoundError:
     __version__ = "0.0.0"
 
+
+def _installed(distribution):
+    try:
+        return version(distribution)
+    except PackageNotFoundError:
+        return None
+
+
+def _check_warp_install():
+    """Report the two ways a Warp install can be wrong before the failure gets cryptic.
+
+    Warp comes from either the ``[warp]`` extra (``warp-lang``) or the ``[neon]``
+    extra (``neon_gpu``, which bundles a fork). Both unpack into
+    ``site-packages/warp``, so having both means whichever pip wrote last owns
+    the directory and uninstalling ``warp-lang`` would delete files Neon needs.
+    XLB no longer pulls ``warp-lang`` in as a core dependency, but pip does not
+    remove one that was already there, so warn rather than assume.
+    """
+    try:
+        import warp  # noqa: F401
+    except ModuleNotFoundError as exc:
+        if exc.name != "warp":
+            raise
+        raise ModuleNotFoundError(
+            "XLB requires Warp, which is not installed. Install exactly one of "
+            "'xlb[warp]' (Warp from PyPI) or 'xlb[neon]' (Neon's bundled Warp fork)."
+        ) from exc
+
+    warp_lang, neon_gpu = _installed("warp-lang"), _installed("neon_gpu")
+    if warp_lang and neon_gpu:
+        warnings.warn(
+            f"Both warp-lang {warp_lang} and neon_gpu {neon_gpu} are installed, and they share "
+            f"site-packages/warp; warp {warp.__version__} currently owns it. Run "
+            "'pip uninstall warp-lang' and reinstall 'xlb[neon]' to restore Neon's fork, "
+            "or drop neon_gpu if you meant to use 'xlb[warp]'.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+
+
+_check_warp_install()
+
 # Enum classes
 from xlb.compute_backend import ComputeBackend as ComputeBackend
 from xlb.precision_policy import PrecisionPolicy as PrecisionPolicy, Precision as Precision
diff --git a/xlb/default_config.py b/xlb/default_config.py
index ab6a331a..ab01ac85 100644
--- a/xlb/default_config.py
+++ b/xlb/default_config.py
@@ -6,6 +6,7 @@
 defaults from :class:`DefaultConfig` when explicit arguments are omitted.
 """
 
+import contextlib
 import os
 
 from xlb.compute_backend import ComputeBackend
@@ -28,14 +29,47 @@ class DefaultConfig:
         Active lattice velocity set.
     default_backend : ComputeBackend or None
         Active compute backend.
+    enable_backward : bool or None
+        Whether Warp generates adjoint (backward) kernels; resolved by
+        :func:`init`. ``None`` until initialized, and stays ``None`` on the JAX
+        backend, which does not use Warp.
     """
 
     default_precision_policy = None
     velocity_set = None
     default_backend = None
+    enable_backward = None
 
 
-def _warp_init_and_select_cuda_device():
+_ENABLE_BACKWARD_ENV = "XLB_WARP_ENABLE_BACKWARD"
+_TRUTHY = {"1", "true", "yes", "on"}
+
+
+def _resolve_enable_backward(enable_backward):
+    """Decide whether Warp should generate adjoint (backward) kernels.
+
+    An explicit ``True``/``False`` from :func:`init` wins. When it is ``None``,
+    fall back to ``XLB_WARP_ENABLE_BACKWARD`` and then to off.
+    """
+    if enable_backward is not None:
+        return bool(enable_backward)
+    return os.environ.get(_ENABLE_BACKWARD_ENV, "").strip().lower() in _TRUTHY
+
+
+def _configure_warp_backward_codegen(enable_backward):
+    """Apply the resolved backward-codegen choice to Warp's global config.
+
+    Warp emits an adjoint version of every kernel by default, which roughly
+    doubles codegen and compile time. XLB's LBM solvers are forward-only, and on
+    the Neon backend the legacy ``@wp.func`` patterns fail NVRTC adjoint
+    compilation outright, so XLB leaves adjoints off unless asked.
+    """
+    import warp as wp
+
+    wp.config.enable_backward = enable_backward
+
+
+def _warp_init_and_select_cuda_device(enable_backward):
     """Initialize Warp and pin the default CUDA device for single-GPU XLB runs.
 
     With multiple GPUs, Warp's default device for allocations and launches can
@@ -44,6 +78,10 @@ def _warp_init_and_select_cuda_device():
     """
     import warp as wp
 
+    # Must precede kernel construction: Warp captures enable_backward into a
+    # module's options when that module's first kernel is created.
+    _configure_warp_backward_codegen(enable_backward)
+
     wp.init()  # TODO: Must be removed in the future versions of WARP
     if wp.get_cuda_device_count() == 0:
         return
@@ -51,13 +89,12 @@ def _warp_init_and_select_cuda_device():
     try:
         wp.set_device(choice)
     except Exception:
-        try:
+        # Fall back to the first GPU; if that fails too, leave Warp's own default.
+        with contextlib.suppress(Exception):
             wp.set_device("cuda:0")
-        except Exception:
-            pass
 
 
-def init(velocity_set, default_backend, default_precision_policy):
+def init(velocity_set, default_backend, default_precision_policy, enable_backward=None):
     """Initialize the global XLB configuration.
 
     Must be called before creating any grid, operator, or field.
@@ -70,13 +107,21 @@ def init(velocity_set, default_backend, default_precision_policy):
         Compute backend to use (JAX, WARP, or NEON).
     default_precision_policy : PrecisionPolicy
         Precision policy for compute and storage dtypes.
+    enable_backward : bool, optional
+        Whether Warp generates adjoint (backward) kernels. Pass ``True`` when
+        differentiating through the solver with ``wp.Tape``; leaving it off
+        roughly halves kernel codegen and compile time, and is required on the
+        Neon backend, whose ``@wp.func`` patterns fail NVRTC adjoint
+        compilation. Defaults to the ``XLB_WARP_ENABLE_BACKWARD`` environment
+        variable, or off when that is unset. Ignored on the JAX backend.
     """
     DefaultConfig.velocity_set = velocity_set
     DefaultConfig.default_backend = default_backend
     DefaultConfig.default_precision_policy = default_precision_policy
 
     if default_backend == ComputeBackend.WARP:
-        _warp_init_and_select_cuda_device()
+        DefaultConfig.enable_backward = _resolve_enable_backward(enable_backward)
+        _warp_init_and_select_cuda_device(DefaultConfig.enable_backward)
     elif default_backend == ComputeBackend.NEON:
         import warp as wp
         import neon
@@ -86,10 +131,11 @@ def init(velocity_set, default_backend, default_precision_policy):
         # wp.config.verbose = True
         # wp.verbose_warnings = True
 
-        _warp_init_and_select_cuda_device()
+        DefaultConfig.enable_backward = _resolve_enable_backward(enable_backward)
+        _warp_init_and_select_cuda_device(DefaultConfig.enable_backward)
 
         # It's a good idea to always clear the kernel cache when developing new native or codegen features
-        wp.build.clear_kernel_cache()
+        wp.clear_kernel_cache()
 
         # !!! DO THIS BEFORE DEFINING/USING ANY KERNELS WITH CUSTOM TYPES
         neon.init()
diff --git a/xlb/helper/initializers.py b/xlb/helper/initializers.py
index 3fbc3f1b..7a2eebed 100644
--- a/xlb/helper/initializers.py
+++ b/xlb/helper/initializers.py
@@ -166,7 +166,7 @@ def jax_implementation(self, bc_mask, f_field):
 
     def _construct_warp(self):
         _q = self.velocity_set.q
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
         _u = _u_vec(self.constant_velocity_vector[0], self.constant_velocity_vector[1], self.constant_velocity_vector[2])
         _rho = self.compute_dtype(self.constant_density)
         _w = self.velocity_set.w
diff --git a/xlb/operator/boundary_condition/bc_equilibrium.py b/xlb/operator/boundary_condition/bc_equilibrium.py
index 85ebe92b..d132a514 100644
--- a/xlb/operator/boundary_condition/bc_equilibrium.py
+++ b/xlb/operator/boundary_condition/bc_equilibrium.py
@@ -81,7 +81,7 @@ def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
 
     def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
         _rho = self.compute_dtype(self.rho)
         _u = _u_vec(self.u[0], self.u[1], self.u[2]) if self.velocity_set.d == 3 else _u_vec(self.u[0], self.u[1])
 
diff --git a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
index 4b7f8f0f..78cbe328 100644
--- a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
+++ b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
@@ -59,7 +59,7 @@ def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
         _opp_indices = self.velocity_set.opp_indices
         _q = wp.constant(self.velocity_set.q)
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         # Construct the functional for this BC
         @wp.func
diff --git a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
index c01259e3..af709d89 100644
--- a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
+++ b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
@@ -89,7 +89,7 @@ def __init__(
             if self.compute_backend in [ComputeBackend.WARP, ComputeBackend.NEON]:
                 if self.velocity_set.d == 2:
                     prescribed_value = np.array([prescribed_value[0], prescribed_value[1], 0.0], dtype=np.float64)
-                prescribed_value = wp.vec(3, dtype=self.compute_dtype)(prescribed_value)
+                prescribed_value = wp.types.vector(length=3, dtype=self.compute_dtype)(prescribed_value)
             self.profile = self._create_constant_prescribed_profile(prescribed_value)
 
     def _create_constant_prescribed_profile(self, prescribed_value):
@@ -101,7 +101,7 @@ def prescribed_profile_jax():
 
             return prescribed_profile_jax
 
-        _u_vec = wp.vec(3, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=3, dtype=self.compute_dtype)
 
         @wp.func
         def prescribed_profile_warp(index: Any, time: Any):
diff --git a/xlb/operator/boundary_condition/bc_hybrid.py b/xlb/operator/boundary_condition/bc_hybrid.py
index 62584160..9ca58b1c 100644
--- a/xlb/operator/boundary_condition/bc_hybrid.py
+++ b/xlb/operator/boundary_condition/bc_hybrid.py
@@ -149,7 +149,7 @@ def __init__(
                 prescribed_value = np.array([prescribed_value[0], prescribed_value[1], 0.0], dtype=np.float64)
 
             # create a constant prescribed profile
-            _u_vec = wp.vec(3, dtype=self.compute_dtype)
+            _u_vec = wp.types.vector(length=3, dtype=self.compute_dtype)
             prescribed_value = _u_vec(prescribed_value)
 
             @wp.func
diff --git a/xlb/operator/boundary_condition/bc_zouhe.py b/xlb/operator/boundary_condition/bc_zouhe.py
index 8e5688c5..a18810e9 100644
--- a/xlb/operator/boundary_condition/bc_zouhe.py
+++ b/xlb/operator/boundary_condition/bc_zouhe.py
@@ -147,10 +147,11 @@ def __init__(
 
     def _create_constant_prescribed_profile(self):
         _prescribed_value = self.prescribed_value
+        _value_vec = wp.types.vector(length=1, dtype=self.precision_policy.store_precision.wp_dtype)
 
         @wp.func
         def prescribed_profile_warp(index: wp.vec3i):
-            return wp.vec(_prescribed_value, length=1)
+            return _value_vec(_prescribed_value)
 
         def prescribed_profile_jax():
             return jnp.array(_prescribed_value, dtype=self.precision_policy.store_precision.jax_dtype).reshape(-1, 1)
diff --git a/xlb/operator/boundary_condition/helper_functions_bc.py b/xlb/operator/boundary_condition/helper_functions_bc.py
index be25db1c..786a4455 100644
--- a/xlb/operator/boundary_condition/helper_functions_bc.py
+++ b/xlb/operator/boundary_condition/helper_functions_bc.py
@@ -60,9 +60,9 @@ def __init__(self, velocity_set=None, precision_policy=None, compute_backend=Non
         _c = self.velocity_set.c
         _c_float = self.velocity_set.c_float
         _qi = self.velocity_set.qi
-        _u_vec = wp.vec(_d, dtype=compute_dtype)
-        _f_vec = wp.vec(_q, dtype=compute_dtype)
-        _missing_mask_vec = wp.vec(_q, dtype=wp.uint8)  # TODO fix vec bool
+        _u_vec = wp.types.vector(length=_d, dtype=compute_dtype)
+        _f_vec = wp.types.vector(length=_q, dtype=compute_dtype)
+        _missing_mask_vec = wp.types.vector(length=_q, dtype=wp.uint8)  # TODO fix vec bool
 
         # Define the operator needed for computing equilibrium
         equilibrium = QuadraticEquilibrium(velocity_set, precision_policy, compute_backend)
@@ -412,7 +412,7 @@ def _construct_warp(self):
         _opp_indices = self.velocity_set.opp_indices
         _id = self.boundary_id
         _num_of_aux_data = self.num_of_aux_data
-        _aux_vec = wp.vec(_num_of_aux_data, dtype=self.compute_dtype)
+        _aux_vec = wp.types.vector(length=_num_of_aux_data, dtype=self.compute_dtype)
 
         @wp.func
         def encoder_functional(
diff --git a/xlb/operator/boundary_masker/aabb_close.py b/xlb/operator/boundary_masker/aabb_close.py
index 21ee31f9..0acfb8a8 100644
--- a/xlb/operator/boundary_masker/aabb_close.py
+++ b/xlb/operator/boundary_masker/aabb_close.py
@@ -307,6 +307,13 @@ def warp_implementation(
             points=wp.array(mesh_vertices, dtype=wp.vec3),
             indices=wp.array(mesh_indices, dtype=wp.int32),
         )
+        # Only the integer ``mesh.id`` reaches the kernels, so retain the Mesh to keep
+        # its device arrays (points/indices/BVH) from being freed while launches are
+        # still pending (use-after-free on the deferred Neon path).
+        if not hasattr(self, "_retained_meshes"):
+            self._retained_meshes = []
+        self._retained_meshes.append(mesh)
+
         mesh_id = wp.uint64(mesh.id)
         bc_id = bc.id
 
diff --git a/xlb/operator/boundary_masker/mesh_boundary_masker.py b/xlb/operator/boundary_masker/mesh_boundary_masker.py
index c6fb778d..00b11a0a 100644
--- a/xlb/operator/boundary_masker/mesh_boundary_masker.py
+++ b/xlb/operator/boundary_masker/mesh_boundary_masker.py
@@ -80,7 +80,7 @@ def pre_compute(
             dist1 = wp.dot(normal, corner - verts[0])
             dist2 = wp.dot(normal, wp.vec3f(1.0, 1.0, 1.0) - corner - verts[0])
 
-            edges = wp.transpose(wp.mat33(verts[1] - verts[0], verts[2] - verts[1], verts[0] - verts[2]))
+            edges = wp.transpose(wp.matrix_from_cols(verts[1] - verts[0], verts[2] - verts[1], verts[0] - verts[2]))
             normal_edge0 = wp.mat33f(0.0)
             normal_edge1 = wp.mat33f(0.0)
             dist_edge = wp.mat33f(0.0)
@@ -141,7 +141,7 @@ def mesh_voxel_intersect(mesh_id: wp.uint64, low: wp.vec3):
                 v2 = wp.mesh_eval_position(mesh_id, f, 0.0, 0.0)
                 normal = wp.mesh_eval_face_normal(mesh_id, f)
 
-                v = wp.transpose(wp.mat33f(v0, v1, v2))
+                v = wp.transpose(wp.matrix_from_cols(v0, v1, v2))
 
                 # TODO: run this on triangles in advance
                 dist1, dist2, normal_edge0, normal_edge1, dist_edge = pre_compute(verts=v, normal=normal)
@@ -208,6 +208,14 @@ def _prepare_kernel_inputs(
             points=wp.array(mesh_vertices, dtype=wp.vec3),
             indices=wp.array(mesh_indices, dtype=wp.int32),
         )
+        # Only the integer ``mesh.id`` is passed to the kernels, which does not keep
+        # the Mesh's device arrays (points/indices/BVH) alive. On the Neon backend the
+        # container is NVRTC-compiled seconds after this function returns, giving Python
+        # time to free the Mesh, so the kernel would dereference freed BVH memory.
+        if not hasattr(self, "_retained_meshes"):
+            self._retained_meshes = []
+        self._retained_meshes.append(mesh)
+
         mesh_id = wp.uint64(mesh.id)
         bc_id = bc.id
         return mesh_id, bc_id
diff --git a/xlb/operator/collision/bgk.py b/xlb/operator/collision/bgk.py
index 29331dc7..8e9a166a 100644
--- a/xlb/operator/collision/bgk.py
+++ b/xlb/operator/collision/bgk.py
@@ -34,7 +34,7 @@ def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, omega):
     def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
         _w = self.velocity_set.w
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         # Construct the functional
         @wp.func
diff --git a/xlb/operator/collision/forced_collision.py b/xlb/operator/collision/forced_collision.py
index 2c2cfa07..9fc29f48 100644
--- a/xlb/operator/collision/forced_collision.py
+++ b/xlb/operator/collision/forced_collision.py
@@ -58,8 +58,7 @@ def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, omega):
 
     def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         # Construct the functional
         @wp.func
diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py
index d814a7cf..6e56c84b 100644
--- a/xlb/operator/collision/kbc.py
+++ b/xlb/operator/collision/kbc.py
@@ -179,8 +179,7 @@ def _construct_warp(self):
             raise NotImplementedError("Velocity set not supported for warp backend: {}".format(type(self.velocity_set)))
 
         # Set local constants TODO: This is a hack and should be fixed with warp update
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
         _epsilon = wp.constant(self.compute_dtype(self.epsilon))
 
         @wp.func
diff --git a/xlb/operator/collision/smagorinsky_les_bgk.py b/xlb/operator/collision/smagorinsky_les_bgk.py
index 7921105e..d66f2d78 100644
--- a/xlb/operator/collision/smagorinsky_les_bgk.py
+++ b/xlb/operator/collision/smagorinsky_les_bgk.py
@@ -70,10 +70,9 @@ def _construct_warp(self):
         _d = self.velocity_set.d
         _cc = self.velocity_set.cc
         _smagorinsky_coef = wp.constant(self.compute_dtype(self.smagorinsky_coef))
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
         _pi_dim = self.velocity_set.d * (self.velocity_set.d + 1) // 2
-        _pi_vec = wp.vec(_pi_dim, dtype=self.compute_dtype)
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _pi_vec = wp.types.vector(length=_pi_dim, dtype=self.compute_dtype)
 
         # Construct the functional
         @wp.func
diff --git a/xlb/operator/equilibrium/multires_quadratic_equilibrium.py b/xlb/operator/equilibrium/multires_quadratic_equilibrium.py
index 58ed96cc..1281d14a 100644
--- a/xlb/operator/equilibrium/multires_quadratic_equilibrium.py
+++ b/xlb/operator/equilibrium/multires_quadratic_equilibrium.py
@@ -29,7 +29,7 @@ def _construct_neon(self):
         functional, _ = self._construct_warp()
 
         # Set local constants TODO: This is a hack and should be fixed with warp update
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
 
         @neon.Container.factory(name="QuadraticEquilibrium")
         def container(
diff --git a/xlb/operator/equilibrium/quadratic_equilibrium.py b/xlb/operator/equilibrium/quadratic_equilibrium.py
index 30ae6dee..6f5b2e07 100644
--- a/xlb/operator/equilibrium/quadratic_equilibrium.py
+++ b/xlb/operator/equilibrium/quadratic_equilibrium.py
@@ -33,8 +33,8 @@ def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
         _c = self.velocity_set.c
         _w = self.velocity_set.w
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
 
         # Construct the equilibrium functional
         @wp.func
@@ -109,7 +109,7 @@ def _construct_neon(self):
         functional, _ = self._construct_warp()
 
         # Set local constants TODO: This is a hack and should be fixed with warp update
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
 
         @neon.Container.factory(name="QuadraticEquilibrium")
         def container(
diff --git a/xlb/operator/force/exact_difference_force.py b/xlb/operator/force/exact_difference_force.py
index 760a9ef1..56fcc8c4 100644
--- a/xlb/operator/force/exact_difference_force.py
+++ b/xlb/operator/force/exact_difference_force.py
@@ -77,7 +77,7 @@ def jax_implementation(self, f_postcollision, feq, rho, u):
 
     def _construct_warp(self):
         _d = self.velocity_set.d
-        _u_vec = wp.vec(_d, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=_d, dtype=self.compute_dtype)
         if _d == 2:
             _force = _u_vec(self.force_vector[0], self.force_vector[1])
         else:
diff --git a/xlb/operator/force/momentum_transfer.py b/xlb/operator/force/momentum_transfer.py
index 9032abb1..38788c6e 100644
--- a/xlb/operator/force/momentum_transfer.py
+++ b/xlb/operator/force/momentum_transfer.py
@@ -70,7 +70,7 @@ def jax_implementation(self, f_0, f_1, bc_mask, missing_mask):
         return f_post_collision, f_post_stream
 
     def _construct_warp(self):
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         @wp.func
         def functional_stream_then_collide(
@@ -168,7 +168,7 @@ def __init__(
 
         if self.compute_backend != ComputeBackend.JAX:
             # Allocate the force vector (the total integral value will be computed)
-            _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+            _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
             self.force = wp.zeros((1), dtype=_u_vec)
 
     @Operator.register_backend(ComputeBackend.JAX)
@@ -214,8 +214,8 @@ def _construct_warp(self):
         # Set local constants
         _c = self.velocity_set.c
         _opp_indices = self.velocity_set.opp_indices
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
-        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.types.vector(length=self.velocity_set.q, dtype=wp.uint8)
         _no_slip_id = self.no_slip_bc_instance.id
 
         # Find velocity index for (0, 0, 0)
diff --git a/xlb/operator/macroscopic/first_moment.py b/xlb/operator/macroscopic/first_moment.py
index 626767fd..d9952ba8 100644
--- a/xlb/operator/macroscopic/first_moment.py
+++ b/xlb/operator/macroscopic/first_moment.py
@@ -19,8 +19,8 @@ def jax_implementation(self, f, rho):
 
     def _construct_warp(self):
         _c = self.velocity_set.c
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.types.vector(length=self.velocity_set.d, dtype=self.compute_dtype)
 
         @wp.func
         def neumaier_sum_component(d: int, f: _f_vec):
diff --git a/xlb/operator/macroscopic/macroscopic.py b/xlb/operator/macroscopic/macroscopic.py
index cb477b9a..5bf4913f 100644
--- a/xlb/operator/macroscopic/macroscopic.py
+++ b/xlb/operator/macroscopic/macroscopic.py
@@ -26,7 +26,7 @@ def jax_implementation(self, f, rho=None, u=None):
         return rho, u
 
     def _construct_warp(self):
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         @wp.func
         def functional(f: _f_vec):
@@ -73,7 +73,7 @@ def _construct_neon(self):
         functional, _ = self._construct_warp()
 
         # Set local vectors
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         @neon.Container.factory("macroscopic")
         def container(
diff --git a/xlb/operator/macroscopic/multires_macroscopic.py b/xlb/operator/macroscopic/multires_macroscopic.py
index abd8d6f2..466bc90f 100644
--- a/xlb/operator/macroscopic/multires_macroscopic.py
+++ b/xlb/operator/macroscopic/multires_macroscopic.py
@@ -34,7 +34,7 @@ def _construct_neon(self):
         # This is because the neon backend relies on the warp functionals for its operations.
         self.zero_moment = ZeroMoment(compute_backend=ComputeBackend.WARP)
         self.first_moment = FirstMoment(compute_backend=ComputeBackend.WARP)
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
         functional, _ = self._construct_warp()
 
         @neon.Container.factory("macroscopic")
diff --git a/xlb/operator/macroscopic/second_moment.py b/xlb/operator/macroscopic/second_moment.py
index 1a0a0f07..79c7ecfd 100644
--- a/xlb/operator/macroscopic/second_moment.py
+++ b/xlb/operator/macroscopic/second_moment.py
@@ -57,12 +57,9 @@ def jax_implementation(
     def _construct_warp(self):
         # Make constants for warp
         _cc = self.velocity_set.cc
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
         _pi_dim = self.velocity_set.d * (self.velocity_set.d + 1) // 2
-        _pi_vec = wp.vec(
-            _pi_dim,
-            dtype=self.compute_dtype,
-        )
+        _pi_vec = wp.types.vector(length=_pi_dim, dtype=self.compute_dtype)
 
         # Construct functional for computing second moment
         @wp.func
diff --git a/xlb/operator/macroscopic/zero_moment.py b/xlb/operator/macroscopic/zero_moment.py
index f536f8d7..eb546878 100644
--- a/xlb/operator/macroscopic/zero_moment.py
+++ b/xlb/operator/macroscopic/zero_moment.py
@@ -17,7 +17,7 @@ def jax_implementation(self, f):
         return jnp.sum(f, axis=0, keepdims=True)
 
     def _construct_warp(self):
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         @wp.func
         def neumaier_sum(f: _f_vec):
diff --git a/xlb/operator/precision_caster/precision_caster.py b/xlb/operator/precision_caster/precision_caster.py
index 5427cba9..6021e3dd 100644
--- a/xlb/operator/precision_caster/precision_caster.py
+++ b/xlb/operator/precision_caster/precision_caster.py
@@ -43,8 +43,8 @@ def jax_implementation(self, f: jnp.ndarray) -> jnp.ndarray:
 
     def _construct_warp(self):
         # Construct needed types and constants
-        from_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.input_precision)
-        to_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.output_precision)
+        from_lattice_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.input_precision)
+        to_lattice_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.output_precision)
         from_array_type = wp.array4d(dtype=self.input_precision)
         to_array_type = wp.array4d(dtype=self.output_precision)
         _q = wp.constant(self.velocity_set.q)
diff --git a/xlb/operator/stepper/ibm_stepper.py b/xlb/operator/stepper/ibm_stepper.py
index 8507b7bf..b7dd4e09 100644
--- a/xlb/operator/stepper/ibm_stepper.py
+++ b/xlb/operator/stepper/ibm_stepper.py
@@ -121,8 +121,7 @@ def jax_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep):
 
     def _construct_ibm_warp(self):
         # Set local constants
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
         _opp_indices = self.velocity_set.opp_indices
         _weights = self.velocity_set.w
         _c = self.velocity_set.c
diff --git a/xlb/operator/stepper/nse_multires_stepper.py b/xlb/operator/stepper/nse_multires_stepper.py
index a9ee8b30..d0141637 100644
--- a/xlb/operator/stepper/nse_multires_stepper.py
+++ b/xlb/operator/stepper/nse_multires_stepper.py
@@ -391,8 +391,8 @@ def _construct_neon(self):
         # in a plain assignment (e.g. `_c = self.velocity_set.c`).  Capturing here
         # makes these values available as simple closure variables.
         lattice_central_index = self.velocity_set.center_index
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.types.vector(length=self.velocity_set.q, dtype=wp.uint8)
         _opp_indices = self.velocity_set.opp_indices
         _c = self.velocity_set.c
 
diff --git a/xlb/operator/stepper/nse_stepper.py b/xlb/operator/stepper/nse_stepper.py
index ee69750e..3c1bbc39 100644
--- a/xlb/operator/stepper/nse_stepper.py
+++ b/xlb/operator/stepper/nse_stepper.py
@@ -334,8 +334,8 @@ def jax_implementation_push(self, f_0, f_1, bc_mask, missing_mask, omega, timest
 
     def _construct_warp(self):
         # Set local constants
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.types.vector(length=self.velocity_set.q, dtype=wp.uint8)
         _opp_indices = self.velocity_set.opp_indices
         lattice_central_index = self.velocity_set.center_index
 
@@ -479,8 +479,8 @@ def _construct_neon(self):
         import neon
 
         # Set local constants
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
-        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.types.vector(length=self.velocity_set.q, dtype=wp.uint8)
         _opp_indices = self.velocity_set.opp_indices
         lattice_central_index = self.velocity_set.center_index
 
diff --git a/xlb/operator/stream/stream.py b/xlb/operator/stream/stream.py
index e337ff39..c7426564 100644
--- a/xlb/operator/stream/stream.py
+++ b/xlb/operator/stream/stream.py
@@ -64,7 +64,7 @@ def _streaming_jax_i(f, c):
     def _construct_warp(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
         _c = self.velocity_set.c
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         # Construct the funcional to get streamed indices
         @wp.func
@@ -127,7 +127,7 @@ def warp_implementation(self, f_0, f_1):
     def _construct_neon(self):
         # Set local constants TODO: This is a hack and should be fixed with warp update
         _c = self.velocity_set.c
-        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _f_vec = wp.types.vector(length=self.velocity_set.q, dtype=self.compute_dtype)
 
         # Construct the funcional to get streamed indices
         @wp.func
diff --git a/xlb/velocity_set/velocity_set.py b/xlb/velocity_set/velocity_set.py
index 6f9634b2..438d7f17 100644
--- a/xlb/velocity_set/velocity_set.py
+++ b/xlb/velocity_set/velocity_set.py
@@ -87,12 +87,12 @@ def _init_warp_properties(self):
         Convert NumPy properties to Warp-specific properties.
         """
         dtype = self.precision_policy.compute_precision.wp_dtype
-        self.c = wp.constant(wp.mat((self.d, self.q), dtype=wp.int32)(self._c))
-        self.w = wp.constant(wp.vec(self.q, dtype=dtype)(self._w))
-        self.opp_indices = wp.constant(wp.vec(self.q, dtype=wp.int32)(self._opp_indices))
-        self.cc = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._cc))
-        self.c_float = wp.constant(wp.mat((self.d, self.q), dtype=dtype)(self._c_float))
-        self.qi = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._qi))
+        self.c = wp.constant(wp.types.matrix(shape=(self.d, self.q), dtype=wp.int32)(self._c))
+        self.w = wp.constant(wp.types.vector(length=self.q, dtype=dtype)(self._w))
+        self.opp_indices = wp.constant(wp.types.vector(length=self.q, dtype=wp.int32)(self._opp_indices))
+        self.cc = wp.constant(wp.types.matrix(shape=(self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._cc))
+        self.c_float = wp.constant(wp.types.matrix(shape=(self.d, self.q), dtype=dtype)(self._c_float))
+        self.qi = wp.constant(wp.types.matrix(shape=(self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._qi))
 
     def _init_neon_properties(self):
         """
@@ -128,13 +128,13 @@ def _init_backend_constants(self):
             self.inv_cs2 = jnp.array(self.inv_cs2, dtype=dtype)
 
     def warp_lattice_vec(self, dtype):
-        return wp.vec(len(self.c), dtype=dtype)
+        return wp.types.vector(length=len(self.c), dtype=dtype)
 
     def warp_u_vec(self, dtype):
-        return wp.vec(self.d, dtype=dtype)
+        return wp.types.vector(length=self.d, dtype=dtype)
 
     def warp_stream_mat(self, dtype):
-        return wp.mat((self.q, self.d), dtype=dtype)
+        return wp.types.matrix(shape=(self.q, self.d), dtype=dtype)
 
     def _construct_qi(self):
         # Qi = cc - cs^2*I