Repository: hpc4cmb/toast (language: python)

Issue:
clean up jax device assignement code
includes:
- small fixes to run with the latest versions of Numpy and JAX
- a clean up of the jax device asignement code


Discussion / hints:

Fix patch:
diff --git a/platforms/linux-nvhpc.sh b/platforms/linux-nvhpc.sh
index c868e5f0d..c2e9e3bcc 100755
--- a/platforms/linux-nvhpc.sh
+++ b/platforms/linux-nvhpc.sh
@@ -12,6 +12,7 @@ cmake \
     -DCMAKE_CXX_FLAGS="-O3 -g -fPIC -pthread -std=c++11" \
     -DOPENMP_TARGET_FLAGS="-Minfo=mp" \
     -DPYTHON_EXECUTABLE:FILEPATH=$(which python3) \
+    -DUSE_OPENMP_TARGET=TRUE \
     -DCMAKE_VERBOSE_MAKEFILE:BOOL=ON \
     ${opts} \
     ..
diff --git a/src/toast/CMakeLists.txt b/src/toast/CMakeLists.txt
index 02b3ac381..e462806c3 100644
--- a/src/toast/CMakeLists.txt
+++ b/src/toast/CMakeLists.txt
@@ -149,7 +149,6 @@ install(FILES
     schedule_sim_ground.py
     schedule_sim_satellite.py
     widgets.py
-    accelerator.py
     "RELEASE"
     DESTINATION ${PYTHON_SITE}/toast
 )
@@ -163,7 +162,9 @@ install(DIRECTORY
 
 # Process the sub directories
 add_subdirectory(io)
+add_subdirectory(accelerator)
 add_subdirectory(tests)
+add_subdirectory(jax)
 add_subdirectory(ops)
 add_subdirectory(templates)
 add_subdirectory(scripts)
diff --git a/src/toast/_libtoast/map_cov.cpp b/src/toast/_libtoast/map_cov.cpp
index 52e3342c1..efb71457a 100644
--- a/src/toast/_libtoast/map_cov.cpp
+++ b/src/toast/_libtoast/map_cov.cpp
@@ -85,7 +85,7 @@ void init_map_cov(py::module & m) {
 
     m.def("cov_accum_diag_hits",
           [](int64_t nsub, int64_t nsubpix, int64_t nnz, py::buffer submap,
-             py::buffer subpix, py::buffer hits) {
+             py::buffer subpix, py::buffer hits, bool use_accel) {
               auto & gt = toast::GlobalTimers::get();
               gt.start("cov_accum_diag_hits");
               pybuffer_check_1D <int64_t> (submap);
@@ -111,7 +111,8 @@ void init_map_cov(py::module & m) {
               return;
           }, py::arg("nsub"), py::arg("nsubpix"), py::arg("nnz"), py::arg("submap"),
           py::arg("subpix"), py::arg(
-              "hits"), R"(
+              "hits"), py::arg(
+              "use_accel"), R"(
         Accumulate hit map.
 
         This uses a pointing matrix to accumulate the local pieces of the hit map.
@@ -133,7 +134,8 @@ void init_map_cov(py::module & m) {
 
     m.def("cov_accum_diag_invnpp",
           [](int64_t nsub, int64_t nsubpix, int64_t nnz, py::buffer submap,
-             py::buffer subpix, py::buffer weights, double scale, py::buffer invnpp) {
+             py::buffer subpix, py::buffer weights, double scale, py::buffer invnpp,
+             bool use_accel) {
               auto & gt = toast::GlobalTimers::get();
               gt.start("cov_accum_diag_invnpp");
               pybuffer_check_1D <int64_t> (submap);
@@ -165,7 +167,8 @@ void init_map_cov(py::module & m) {
               return;
           }, py::arg("nsub"), py::arg("nsubpix"), py::arg("nnz"), py::arg("submap"),
           py::arg("subpix"), py::arg("weights"), py::arg("scale"), py::arg(
-              "invnpp"), R"(
+              "invnpp"), py::arg(
+              "use_accel"), R"(
         Accumulate block diagonal noise covariance.
 
         This uses a pointing matrix to accumulate the local pieces
diff --git a/src/toast/_libtoast/ops_mapmaker_utils.cpp b/src/toast/_libtoast/ops_mapmaker_utils.cpp
index cbf239163..06a066334 100644
--- a/src/toast/_libtoast/ops_mapmaker_utils.cpp
+++ b/src/toast/_libtoast/ops_mapmaker_utils.cpp
@@ -34,8 +34,7 @@ void build_noise_weighted_inner(
     uint8_t shared_mask,
     int64_t n_pix_submap,
     bool use_shared_flags,
-    bool use_det_flags
-) {
+    bool use_det_flags) {
     int32_t w_indx = weight_index[idet];
     int32_t p_indx = pixel_index[idet];
     int32_t f_indx = flag_index[idet];
@@ -63,8 +62,7 @@ void build_noise_weighted_inner(
     if (
         (pixels[off_p] >= 0) &&
         (det_check == 0) &&
-        (shared_check == 0)
-    ) {
+        (shared_check == 0)) {
         // Good data, accumulate
         global_submap = (int64_t)(pixels[off_p] / n_pix_submap);
 
@@ -111,8 +109,8 @@ void init_ops_mapmaker_utils(py::module & m) {
             py::buffer intervals,
             py::buffer shared_flags,
             uint8_t shared_flag_mask,
-            bool use_accel
-        ) {
+            bool use_accel)
+        {
             auto & omgr = OmpManager::get();
             int dev = omgr.get_device();
             bool offload = (!omgr.device_is_host()) && use_accel;
@@ -200,33 +198,32 @@ void init_ops_mapmaker_utils(py::module & m) {
                 uint8_t * dev_shared_flags = omgr.device_ptr(raw_shared_flags);
                 uint8_t * dev_det_flags = omgr.device_ptr(raw_det_flags);
 
-                # pragma omp target data             \
-                device(dev)                          \
-                map(to:                              \
-                raw_weight_index[0:n_det],           \
-                raw_pixel_index[0:n_det],            \
-                raw_flag_index[0:n_det],             \
-                raw_data_index[0:n_det],             \
-                raw_det_scale[0:n_det],              \
-                raw_global2local[0:n_global_submap], \
-                n_view,                              \
-                n_det,                               \
-                n_samp,                              \
-                nnz,                                 \
-                n_pix_submap,                        \
-                det_flag_mask,                       \
-                shared_flag_mask,                    \
-                use_shared_flags,                    \
-                use_det_flags                        \
-                )                                    \
-                use_device_ptr(                      \
-                raw_weight_index,                    \
-                raw_pixel_index,                     \
-                raw_flag_index,                      \
-                raw_data_index,                      \
-                raw_det_scale,                       \
-                raw_global2local                     \
-                )
+                # pragma omp target data              \
+                device(dev)                           \
+                map(to                                \
+                :                                     \
+                raw_weight_index [0:n_det],           \
+                raw_pixel_index [0:n_det],            \
+                raw_flag_index [0:n_det],             \
+                raw_data_index [0:n_det],             \
+                raw_det_scale [0:n_det],              \
+                raw_global2local [0:n_global_submap], \
+                n_view,                               \
+                n_det,                                \
+                n_samp,                               \
+                nnz,                                  \
+                n_pix_submap,                         \
+                det_flag_mask,                        \
+                shared_flag_mask,                     \
+                use_shared_flags,                     \
+                use_det_flags)                        \
+                use_device_ptr(                       \
+                raw_weight_index,                     \
+                raw_pixel_index,                      \
+                raw_flag_index,                       \
+                raw_data_index,                       \
+                raw_det_scale,                        \
+                raw_global2local)
                 {
                     # pragma omp target teams distribute collapse(2) \
                     is_device_ptr(                                   \
@@ -236,8 +233,7 @@ void init_ops_mapmaker_utils(py::module & m) {
                     dev_det_flags,                                   \
                     dev_intervals,                                   \
                     dev_shared_flags,                                \
-                    dev_zmap                                         \
-                    )
+                    dev_zmap)
                     for (int64_t idet = 0; idet < n_det; idet++) {
                         for (int64_t iview = 0; iview < n_view; iview++) {
                             # pragma omp parallel for default(shared)
@@ -317,9 +313,11 @@ void init_ops_mapmaker_utils(py::module & m) {
                                 use_shared_flags,
                                 use_det_flags
                             );
-                            #pragma omp critical
+
+                            // #pragma omp critical
                             {
                                 for (int64_t iw = 0; iw < nnz; iw++) {
+                                    #pragma omp atomic
                                     raw_zmap[zoff + iw] += zmap_val[iw];
                                 }
                             }
diff --git a/src/toast/_libtoast/tod_filter.cpp b/src/toast/_libtoast/tod_filter.cpp
index 5b80c130f..82cecc800 100644
--- a/src/toast/_libtoast/tod_filter.cpp
+++ b/src/toast/_libtoast/tod_filter.cpp
@@ -319,7 +319,7 @@ void init_tod_filter(py::module & m) {
 
     m.def("filter_polynomial",
           [](int64_t order, py::buffer flags, py::list signals, py::buffer starts,
-             py::buffer stops) {
+             py::buffer stops, bool use_accel) {
               pybuffer_check_1D <uint8_t> (flags);
               pybuffer_check_1D <int64_t> (starts);
               pybuffer_check_1D <int64_t> (stops);
@@ -359,7 +359,8 @@ void init_tod_filter(py::module & m) {
               return;
           }, py::arg("order"), py::arg("flags"), py::arg("signals"), py::arg("starts"),
           py::arg(
-              "stops"), R"(
+              "stops"), py::arg(
+              "use_accel"), R"(
         Fit and subtract a polynomial from one or more signals.
 
         Args:
@@ -376,7 +377,7 @@ void init_tod_filter(py::module & m) {
 
     m.def("filter_poly2D",
           [](py::buffer det_groups, py::buffer templates, py::buffer signals,
-             py::buffer masks, py::buffer coeff) {
+             py::buffer masks, py::buffer coeff, bool use_accel) {
               pybuffer_check <uint8_t> (masks);
               pybuffer_check <int32_t> (det_groups);
               pybuffer_check <double> (templates);
@@ -465,7 +466,8 @@ void init_tod_filter(py::module & m) {
           }, py::arg("det_groups"), py::arg("templates"), py::arg("signals"),
           py::arg("masks"),
           py::arg(
-              "coeff"), R"(
+              "coeff"), py::arg(
+              "use_accel"), R"(
         Solves for 2D polynomial coefficients at each sample.
 
         Args:
diff --git a/src/toast/_libtoast/todmap_scanning.cpp b/src/toast/_libtoast/todmap_scanning.cpp
index 589403e83..b68dd3d50 100644
--- a/src/toast/_libtoast/todmap_scanning.cpp
+++ b/src/toast/_libtoast/todmap_scanning.cpp
@@ -27,6 +27,8 @@ void register_scan_map(py::module & m, char const * name) {
                   auto log = toast::Logger::get();
                   std::ostringstream o;
                   o << "Buffer sizes are not consistent.";
+                  o << " submap array (" << info_submap.size;
+                  o << ") != tod nsamp (" << nsamp << ")";
                   log.error(o.str().c_str());
                   throw std::runtime_error(o.str().c_str());
               }
@@ -35,6 +37,8 @@ void register_scan_map(py::module & m, char const * name) {
                   auto log = toast::Logger::get();
                   std::ostringstream o;
                   o << "Buffer sizes are not consistent.";
+                  o << " weights (" << info_weights.size << " / " << nmap;
+                  o << ") != tod nsamp (" << nsamp << ")";
                   log.error(o.str().c_str());
                   throw std::runtime_error(o.str().c_str());
               }
diff --git a/src/toast/accelerator/CMakeLists.txt b/src/toast/accelerator/CMakeLists.txt
new file mode 100644
index 000000000..2d87e4ffb
--- /dev/null
+++ b/src/toast/accelerator/CMakeLists.txt
@@ -0,0 +1,10 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    accel.py
+    data_localization.py
+    kernel_registry.py
+    DESTINATION ${PYTHON_SITE}/toast/accelerator
+)
diff --git a/src/toast/accelerator/__init__.py b/src/toast/accelerator/__init__.py
new file mode 100644
index 000000000..0af03c270
--- /dev/null
+++ b/src/toast/accelerator/__init__.py
@@ -0,0 +1,20 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import os
+
+from .accel import (
+    AcceleratorObject,
+    accel_assign_device,
+    accel_data_create,
+    accel_data_delete,
+    accel_data_present,
+    accel_data_update_device,
+    accel_data_update_host,
+    accel_enabled,
+    accel_get_device,
+    use_accel_jax,
+    use_accel_omp,
+)
+from .kernel_registry import ImplementationType, kernel
diff --git a/src/toast/accelerator.py b/src/toast/accelerator/accel.py
similarity index 67%
rename from src/toast/accelerator.py
rename to src/toast/accelerator/accel.py
index 1b1867599..261d6e882 100644
--- a/src/toast/accelerator.py
+++ b/src/toast/accelerator/accel.py
@@ -1,18 +1,18 @@
-# Copyright (c) 2015-2022 by the parties listed in the AUTHORS file.
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
 # All rights reserved.  Use of this source code is governed by
 # a BSD-style license that can be found in the LICENSE file.
 
 import os
-import time
 
-from ._libtoast import Logger
-from ._libtoast import accel_create as omp_accel_create
-from ._libtoast import accel_delete as omp_accel_delete
-from ._libtoast import accel_enabled as omp_accel_enabled
-from ._libtoast import accel_get_device as omp_accel_get_device
-from ._libtoast import accel_present as omp_accel_present
-from ._libtoast import accel_update_device as omp_accel_update_device
-from ._libtoast import accel_update_host as omp_accel_update_host
+from .._libtoast import Logger
+from .._libtoast import accel_assign_device as omp_accel_assign_device
+from .._libtoast import accel_create as omp_accel_create
+from .._libtoast import accel_delete as omp_accel_delete
+from .._libtoast import accel_enabled as omp_accel_enabled
+from .._libtoast import accel_get_device as omp_accel_get_device
+from .._libtoast import accel_present as omp_accel_present
+from .._libtoast import accel_update_device as omp_accel_update_device
+from .._libtoast import accel_update_host as omp_accel_update_host
 
 enable_vals = ["1", "yes", "true"]
 
@@ -23,22 +23,27 @@
     else:
         log = Logger.get()
         msg = "TOAST_GPU_OPENMP enabled at runtime, but package was not built "
-        msg += "with OpenMP target offload support.  Disabling."
-        log.warning(msg)
+        msg += "with OpenMP target offload support."
+        log.error(msg)
+        raise RuntimeError(msg)
 
 use_accel_jax = False
-if "TOAST_GPU_JAX" in os.environ and os.environ["TOAST_GPU_JAX"] in enable_vals:
+if ("TOAST_GPU_JAX" in os.environ) and (os.environ["TOAST_GPU_JAX"] in enable_vals):
     try:
         import jax
-        import jax.numpy as jnp
+
+        from ..jax.device import jax_accel_assign_device, jax_accel_get_device
+        from ..jax.intervals import INTERVALS_JAX
+        from ..jax.mutableArray import MutableJaxArray
 
         use_accel_jax = True
-    except Exception:
+    except ImportError:
         # There could be many possible exceptions...
         log = Logger.get()
         msg = "TOAST_GPU_JAX enabled at runtime, but jax is not "
-        msg += "importable.  Disabling."
-        log.warning(msg)
+        msg += "importable."
+        log.error(msg)
+        raise RuntimeError(msg)
 
 if use_accel_omp and use_accel_jax:
     log = Logger.get()
@@ -46,12 +51,6 @@
     log.error(msg)
     raise RuntimeError(msg)
 
-
-# This stores the assigned jax device for this process, computed during
-# MPI initialization.
-jax_local_device = None
-
-
 # Wrapper functions that work with either numpy arrays mapped to omp device memory
 # or jax arrays.
 
@@ -63,25 +62,45 @@ def accel_enabled():
 
 def accel_get_device():
     """Return the device ID assigned to this process."""
-    global jax_local_device
     if use_accel_omp:
         return omp_accel_get_device()
     elif use_accel_jax:
-        if jax_local_device is None:
-            raise RuntimeError("Jax device is not set!")
-        return jax_local_device
+        return jax_accel_get_device()
     else:
         log = Logger.get()
         log.warning("Accelerator support not enabled, returning device -1")
         return -1
 
 
+def accel_assign_device(node_procs, node_rank, disabled):
+    """
+    Assign processes to target devices.
+
+    NOTE:
+    One can pick devices visible to processes using Slurm and teh following commands
+    `--gpus-per-task=1 --gpu-bind=single:1`
+
+    Args:
+        node_procs (int): number of processes per node
+        node_rank (int): rank of the current process, within the node
+        disabled (bool): gpu computing is disabled
+
+    Returns:
+        None: the device is stored in a backend specific global variable
+    """
+    # FIXME some functions (such as poiting_detector) require the omp device to have been assigned
+    # so it should be called even when using JAX or running on CPU
+    omp_accel_assign_device(node_procs, node_rank, disabled)
+    if use_accel_jax:
+        jax_accel_assign_device(node_procs, node_rank, disabled)
+
+
 def accel_data_present(data):
     """Check if data is present on the device.
 
     For OpenMP target offload, this checks if the input data has an entry in the
-    global map of host to device pointers.  For jax, this tests if the input array
-    is a jax array.
+    global map of host to device pointers.
+    For jax, this tests if the input array is a jax array.
 
     Args:
         data (array):  The data to test.
@@ -90,17 +109,19 @@ def accel_data_present(data):
         (bool):  True if the data is present on the device.
 
     """
+    log = Logger.get()
+    log.verbose("accel_data_present")
     if data is None:
         return False
-    if use_accel_omp:
+    elif use_accel_omp:
         return omp_accel_present(data)
     elif use_accel_jax:
-        if isinstance(data, jnp.DeviceArray):
-            return True
-        else:
-            return False
+        return (
+            isinstance(data, MutableJaxArray)
+            or isinstance(data, jax.numpy.ndarray)
+            or isinstance(data, INTERVALS_JAX)
+        )
     else:
-        log = Logger.get()
         log.warning("Accelerator support not enabled, data not present")
         return False
 
@@ -110,20 +131,19 @@ def accel_data_create(data):
 
     Using the input data array, create a corresponding device array.  For OpenMP
     target offload, this allocates device memory and adds it to the global map
-    of host to device pointers.  For jax arrays, this is a no-op, since those
-    arrays are mapped and managed elsewhere.
+    of host to device pointers. For jax it just wraps the numpy array
 
     Args:
         data (array):  The host array.
 
     Returns:
-        None
+        None for OpenMP target offload and a JAX array for JAX
 
     """
     if use_accel_omp:
         omp_accel_create(data)
     elif use_accel_jax:
-        pass
+        return MutableJaxArray(data)
     else:
         log = Logger.get()
         log.warning("Accelerator support not enabled, cannot create")
@@ -140,14 +160,14 @@ def accel_data_update_device(data):
         data (array):  The host array.
 
     Returns:
-        (object):  Either the original input (for OpenMP) or a jax array.
+        (object):  Either the original input (for OpenMP) or a new jax array.
 
     """
     if use_accel_omp:
         omp_accel_update_device(data)
         return data
     elif use_accel_jax:
-        return jnp.DeviceArray(data)
+        return MutableJaxArray(data)
     else:
         log = Logger.get()
         log.warning("Accelerator support not enabled, not updating device")
@@ -173,12 +193,7 @@ def accel_data_update_host(data):
         omp_accel_update_host(data)
         return data
     elif use_accel_jax:
-        if isinstance(data, jnp.DeviceArray):
-            # Return a numpy array
-            return data.copy()
-        else:
-            # Already on the host
-            return data
+        return data.to_host()
     else:
         log = Logger.get()
         log.warning("Accelerator support not enabled, not updating host")
@@ -188,7 +203,9 @@ def accel_data_delete(data):
     """Delete device copy of the data.
 
     For OpenMP target offload, this deletes the device allocated memory and removes
-    the host entry from the global memory map.  For jax, this is a no-op.
+    the host entry from the global memory map.
+
+    For jax, this returns a host array (if needed).
 
     Args:
         data (array):  The host array.
@@ -200,10 +217,13 @@ def accel_data_delete(data):
     if use_accel_omp:
         omp_accel_delete(data)
     elif use_accel_jax:
-        pass
+        # if needed, make sure that data is on host
+        if accel_data_present(data):
+            data = data.host_data
     else:
         log = Logger.get()
         log.warning("Accelerator support not enabled, cannot delete device data")
+    return data
 
 
 class AcceleratorObject(object):
@@ -268,14 +288,14 @@ def accel_used(self, state):
         self._accel_used = state
 
     def _accel_create(self):
-        pass
+        msg = f"The _accel_create function was not defined for this class."
+        raise RuntimeError(msg)
 
     def accel_create(self):
-        """Create a copy of the data on the accelerator.
+        """Create a (potentially uninitialized) copy of the data on the accelerator.
 
         Returns:
             None
-
         """
         if not accel_enabled():
             return
@@ -287,24 +307,26 @@ def accel_create(self):
         self._accel_create()
 
     def _accel_update_device(self):
-        pass
+        msg = f"The _accel_update_device function was not defined for this class."
+        raise RuntimeError(msg)
 
     def accel_update_device(self):
         """Copy the data to the accelerator.
 
         Returns:
             None
-
         """
         if not accel_enabled():
             return
-        if not self.accel_exists():
+        if (not self.accel_exists()) and (not use_accel_jax):
+            # There is no data on device
+            # NOTE: this does no apply to JAX as JAX will allocate on the fly
             log = Logger.get()
             msg = f"Data does not exist on device, cannot update"
             log.error(msg)
             raise RuntimeError(msg)
         if self.accel_in_use():
-            # The active copy is on the device
+            # The active copy is already on the device
             log = Logger.get()
             msg = f"Active data is already on device, cannot update"
             log.error(msg)
@@ -313,14 +335,14 @@ def accel_update_device(self):
         self.accel_used(True)
 
     def _accel_update_host(self):
-        pass
+        msg = f"The _accel_update_host function was not defined for this class."
+        raise RuntimeError(msg)
 
     def accel_update_host(self):
         """Copy the data to the host from the accelerator.
 
         Returns:
             None
-
         """
         if not accel_enabled():
             return
@@ -330,7 +352,7 @@ def accel_update_host(self):
             log.error(msg)
             raise RuntimeError(msg)
         if not self.accel_in_use():
-            # The active copy is on the host
+            # The active copy is already on the host
             log = Logger.get()
             msg = f"Active data is already on host, cannot update"
             log.error(msg)
@@ -339,7 +361,8 @@ def accel_update_host(self):
         self.accel_used(False)
 
     def _accel_delete(self):
-        pass
+        msg = f"The _accel_delete function was not defined for this class."
+        raise RuntimeError(msg)
 
     def accel_delete(self):
         """Delete the data from the accelerator.
@@ -350,7 +373,8 @@ def accel_delete(self):
         """
         if not accel_enabled():
             return
-        if not self.accel_exists():
+        if (not self.accel_exists()) and (not use_accel_jax):
+            # NOTE: this check does not apply to JAX as the data will not be on device after an update_host
             log = Logger.get()
             msg = f"Data does not exist on device, cannot delete"
             log.error(msg)
diff --git a/src/toast/accelerator/data_localization.py b/src/toast/accelerator/data_localization.py
new file mode 100644
index 000000000..32320cdcc
--- /dev/null
+++ b/src/toast/accelerator/data_localization.py
@@ -0,0 +1,137 @@
+import os
+import sys
+from collections import defaultdict
+from functools import wraps
+
+from ..accelerator import accel_data_present
+from ..utils import Logger
+
+# ---------------------------------------------------------------------------
+# RECORDING CLASS
+
+use_debug_assert = ("TOAST_LOGLEVEL" in os.environ) and (
+    os.environ["TOAST_LOGLEVEL"] in ["DEBUG", "VERBOSE"]
+)
+"""
+Assert is used only if `TOAST_LOGLEVEL` is set to `DEBUG`.
+"""
+
+
+def bytes_of_data(data):
+    """
+    Returns the size of an input in bytes.
+    """
+    if hasattr(data, "nbytes"):
+        return data.nbytes
+    else:
+        return sys.getsizeof(data)
+
+
+class DataMovementRecord:
+    """
+    data structure used to track data movement to and from the GPU for a particular function
+    """
+
+    def __init__(self):
+        self.nb_calls = 0
+        self.input_bytes = 0
+        self.input_to_gpu_bytes = 0
+
+    def add(self, args, kwargs, display=False):
+        """
+        Adds data to the record, returns True if at least one of the inputs was on GPU
+        """
+        self.nb_calls += 1
+        some_gpu_input = False
+        # args
+        for i, input in enumerate(args):
+            bytes = bytes_of_data(input)
+            gpu_data = accel_data_present(input)
+            self.input_bytes += bytes
+            if not gpu_data:
+                self.input_to_gpu_bytes += bytes
+            if display:
+                print(f"input[{i}] size:{bytes} GPU:{gpu_data}")
+            some_gpu_input = some_gpu_input or gpu_data
+        # kwargs
+        for name, input in kwargs.items():
+            bytes = bytes_of_data(input)
+            gpu_data = accel_data_present(input)
+            self.input_bytes += bytes
+            if not gpu_data:
+                self.input_to_gpu_bytes += bytes
+            if display:
+                print(f"input[{name}] size:{bytes} GPU:{gpu_data}")
+            some_gpu_input = some_gpu_input or gpu_data
+        return some_gpu_input
+
+
+class DataMovementTracker:
+    """
+    data structure used to track data movement to and from the GPU for several functions
+    """
+
+    def __init__(self):
+        # each function gets a record
+        self.records = defaultdict(lambda: DataMovementRecord())
+
+    def add(self, functionname, args, kwargs, display=False):
+        """
+        keep track of the total size of the inputs
+        and the size of the inputs already on GPU
+        this is useful to monitor data movement on a per function basis
+        """
+        # records the data size and movement
+        if display:
+            print(f"DATA MOVEMENT ({functionname}):")
+        some_gpu_input = self.records[functionname].add(args, kwargs, display)
+        # send a warning in case of suspicious data movement
+        use_accel = kwargs.get("use_accel", False)
+        if use_accel and (not some_gpu_input):
+            msg = f"function '{functionname}' has NO input on GPU despite use_accel=True!"
+            log = Logger.get()
+            log.warning(msg)
+        elif some_gpu_input and (not use_accel):
+            msg = f"function '{functionname}' has inputs on GPU despite use_accel=False!"
+            log = Logger.get()
+            log.warning(msg)
+
+    def __str__(self):
+        """
+        produces a CSV representation of the DataMovementTracker
+        with one line per function and using ',' as the separator
+        """
+        if not use_debug_assert:
+            return "DataMovementTracker: no monitoring unless TOAST_LOGLEVEL is set to DEBUG or VERBOSE."
+        result = "DataMovementTracker:\n-----\n"
+        result += "function, nb calls, input bytes, input bytes moved to GPU\n"
+        for function_name, record in sorted(self.records.items()):
+            result += f"{function_name}, {record.nb_calls}, {record.input_bytes}, {record.input_to_gpu_bytes}\n"
+        result += "-----"
+        return result
+
+
+# ----------------------------------------------------------------------------------
+# OPERATIONS
+
+dataMovementTracker = DataMovementTracker()
+"""global variable used to track data movement to and from the GPU"""
+
+
+def function_datamovementtracker(f):
+    """
+    Wraps a function to keep track of movement to and from the GPU
+    NOTE: the recording is only done if we are in DEBUG or VERBOSE mode
+    """
+
+    @wraps(f)
+    def f_result(*args, **kwargs):
+        dataMovementTracker.add(f.__name__, args, kwargs)
+        return f(*args, **kwargs)
+
+    return f_result if use_debug_assert else f
+
+
+def display_datamovement():
+    """Displays all recorded data movement so far"""
+    print(dataMovementTracker)
diff --git a/src/toast/accelerator/kernel_registry.py b/src/toast/accelerator/kernel_registry.py
new file mode 100644
index 000000000..c4a4e7ea0
--- /dev/null
+++ b/src/toast/accelerator/kernel_registry.py
@@ -0,0 +1,90 @@
+# Copyright (c) 2023-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import inspect
+from enum import IntEnum
+from functools import wraps
+
+from ..timing import function_timer
+from .accel import accel_enabled
+from .data_localization import function_datamovementtracker
+
+
+class ImplementationType(IntEnum):
+    """Describes the various implementation kinds"""
+
+    DEFAULT = 0
+    COMPILED = 1
+    NUMPY = 2
+    JAX = 3
+
+
+registry = dict()
+
+
+def kernel(impl, name=None):
+    """Decorator which registers a kernel function.
+
+    This will associate the function with a particular implementation of a kernel name.
+    It checks the inputs and ensures that multiple functions are not registered for
+    the same name / implementation.
+
+    This returns the "dispatch" function for the kernel.  Since we have a decorator
+    that takes some arguments, we must use a nested scheme.
+
+    Args:
+        impl (ImplementationType):  The implementation type
+        name (str):  If not None, the kernel name to use for registration.  Otherwise
+            the function name is used
+
+    Returns:
+        (function):  The dispatch function.
+
+    """
+    global registry
+
+    if not isinstance(impl, ImplementationType):
+        raise ValueError(
+            "kernel decorator second argument should be ImplementationType"
+        )
+
+    def _kernel(f):
+        nonlocal impl
+        nonlocal name
+        if name is None:
+            name = f.__name__
+
+        if name not in registry:
+            # This is the first occurence of this kernel.  Set up the dictionary of
+            # implementations and create a dispatch function.
+            registry[name] = dict()
+
+            def dispatch(
+                *args, impl=ImplementationType.DEFAULT, use_accel=False, **kwargs
+            ):
+                return registry[name][impl](*args, use_accel=use_accel, **kwargs)
+
+            registry[name]["dispatch"] = dispatch
+
+        if impl in registry[name]:
+            msg = f"Implementation {impl.name} for kernel {name} already exists."
+            raise RuntimeError(msg)
+
+        if impl == ImplementationType.DEFAULT:
+            # When we register the default implementation, assign its docstring to
+            # the dispatch function
+            registry[name]["dispatch"].__doc__ = f.__doc__
+
+        # Register the function with timer and data movement tracker
+        if (
+            (impl == ImplementationType.JAX) or (impl == ImplementationType.COMPILED)
+        ) and accel_enabled():
+            registry[name][impl] = function_timer(function_datamovementtracker(f))
+        else:
+            registry[name][impl] = function_timer(f)
+
+        # Return the dispatch function in place of the original
+        return registry[name]["dispatch"]
+
+    return _kernel
diff --git a/src/toast/data.py b/src/toast/data.py
index 15f6fabf7..552706441 100644
--- a/src/toast/data.py
+++ b/src/toast/data.py
@@ -8,8 +8,7 @@
 
 import numpy as np
 
-from ._libtoast import accel_enabled
-from .accelerator import AcceleratorObject
+from .accelerator import AcceleratorObject, accel_enabled
 from .mpi import Comm
 from .utils import Logger
 
@@ -454,30 +453,49 @@ def accel_create(self, names):
             None
 
         """
+        log = Logger.get()
         if not accel_enabled():
+            log.verbose(f"accel_enabled is False, canceling accel_create.")
             return
-        log = Logger.get()
+
         for ob in self.obs:
             for key in names["detdata"]:
                 if not ob.detdata.accel_exists(key):
-                    log.verbose(f"Calling ob {ob.name} detdata accel_create for {key}")
+                    log.verbose(f"ob {ob.name} detdata: accel_create '{key}'")
                     ob.detdata.accel_create(key)
+                else:
+                    log.verbose(
+                        f"ob {ob.name} detdata: accel_create '{key}' already on device"
+                    )
             for key in names["shared"]:
                 if not ob.shared.accel_exists(key):
-                    log.verbose(f"Calling ob {ob.name} shared accel_create for {key}")
+                    log.verbose(f"ob {ob.name} shared: accel_create '{key}'")
                     ob.shared.accel_create(key)
+                else:
+                    log.verbose(
+                        f"ob {ob.name} shared: accel_create '{key}' already on device"
+                    )
             for key in names["intervals"]:
                 if not ob.intervals.accel_exists(key):
+                    log.verbose(f"ob {ob.name} intervals: accel_create '{key}'")
+                    ob.intervals.accel_create(key)
+                else:
                     log.verbose(
-                        f"Calling ob {ob.name} intervals accel_create for {key}"
+                        f"ob {ob.name} intervals: accel_create '{key}' already on device"
                     )
-                    ob.intervals.accel_create(key)
+
         for key in names["global"]:
             val = self._internal[key]
             if isinstance(val, AcceleratorObject):
                 if not val.accel_exists():
-                    log.verbose(f"Calling Data accel_create for {key}")
+                    log.verbose(f"Data accel_create: '{key}'")
                     val.accel_create()
+                else:
+                    log.verbose(f"Data accel_create: '{key}' already on device")
+            else:
+                log.verbose(
+                    f"Data accel_create: '{key}' ({type(val)}) is not an AcceleratorObject"
+                )
 
     def accel_update_device(self, names):
         """Copy a set of data objects to the device.
@@ -498,40 +516,48 @@ def accel_update_device(self, names):
         for ob in self.obs:
             for key in names["detdata"]:
                 if ob.detdata.accel_in_use(key):
-                    msg = f"Skipping {ob.name} detdata update_device for {key}, "
+                    msg = f"Skipping {ob.name} detdata update_device for '{key}', "
                     msg += "device data in use"
                     log.verbose(msg)
                 else:
-                    log.verbose(f"Calling ob {ob.name} detdata update_device for {key}")
+                    log.verbose(
+                        f"Calling ob {ob.name} detdata update_device for '{key}'"
+                    )
                     ob.detdata.accel_update_device(key)
             for key in names["shared"]:
                 if ob.shared.accel_in_use(key):
-                    msg = f"Skipping {ob.name} shared update_device for {key}, "
+                    msg = f"Skipping {ob.name} shared update_device for '{key}', "
                     msg += "device data in use"
                     log.verbose(msg)
                 else:
-                    log.verbose(f"Calling ob {ob.name} shared update_device for {key}")
+                    log.verbose(
+                        f"Calling ob {ob.name} shared update_device for '{key}'"
+                    )
                     ob.shared.accel_update_device(key)
             for key in names["intervals"]:
                 if ob.intervals.accel_in_use(key):
-                    msg = f"Skipping {ob.name} intervals update_device for {key}, "
+                    msg = f"Skipping {ob.name} intervals update_device for '{key}', "
                     msg += "device data in use"
                     log.verbose(msg)
                 else:
                     log.verbose(
-                        f"Calling ob {ob.name} intervals update_device for {key}"
+                        f"Calling ob {ob.name} intervals update_device for '{key}'"
                     )
                     ob.intervals.accel_update_device(key)
         for key in names["global"]:
             val = self._internal[key]
             if isinstance(val, AcceleratorObject):
                 if val.accel_in_use():
-                    msg = f"Skipping update_device for {key}, "
+                    msg = f"Skipping update_device for '{key}', "
                     msg += "device data in use"
                     log.verbose(msg)
                 else:
-                    log.verbose(f"Calling Data update_device for {key}")
+                    log.verbose(f"Calling Data update_device for '{key}'")
                     val.accel_update_device()
+            else:
+                msg = f"Data accel_update_device: '{key}' ({type(val)}) "
+                msg += "is not an AcceleratorObject"
+                log.verbose(msg)
 
     def accel_update_host(self, names):
         """Copy a set of data objects to the host.
@@ -553,58 +579,62 @@ def accel_update_host(self, names):
             for key in names["detdata"]:
                 if ob.detdata.accel_exists(key):
                     if not ob.detdata.accel_in_use(key):
-                        msg = f"Skipping {ob.name} detdata update_host for {key}, "
+                        msg = f"Skipping {ob.name} detdata update_host for '{key}', "
                         msg += "host data in use"
                         log.verbose(msg)
                     else:
                         log.verbose(
-                            f"Calling ob {ob.name} detdata update_host for {key}"
+                            f"Calling ob {ob.name} detdata update_host for '{key}'"
                         )
                         ob.detdata.accel_update_host(key)
                 else:
-                    log.verbose(
-                        f"Skip update_host for ob {ob.name} detdata {key}, data not present"
-                    )
+                    msg = f"Skip update_host for ob {ob.name} detdata '{key}', "
+                    msg += "data not present"
+                    log.verbose(msg)
             for key in names["shared"]:
                 if ob.shared.accel_exists(key):
                     if not ob.shared.accel_in_use(key):
-                        msg = f"Skipping {ob.name} shared update_host for {key}, "
+                        msg = f"Skipping {ob.name} shared update_host for '{key}', "
                         msg += "host data in use"
                         log.verbose(msg)
                     else:
                         log.verbose(
-                            f"Calling ob {ob.name} shared update_host for {key}"
+                            f"Calling ob {ob.name} shared update_host for '{key}'"
                         )
                         ob.shared.accel_update_host(key)
                 else:
-                    log.verbose(
-                        f"Skip update_host for ob {ob.name} shared {key}, data not present"
-                    )
+                    msg = f"Skip update_host for ob {ob.name} shared '{key}', "
+                    msg += "data not present"
+                    log.verbose(msg)
             for key in names["intervals"]:
                 if ob.intervals.accel_exists(key):
                     if not ob.intervals.accel_in_use(key):
-                        msg = f"Skipping {ob.name} intervals update_host for {key}, "
+                        msg = f"Skipping {ob.name} intervals update_host for '{key}', "
                         msg += "host data in use"
                         log.verbose(msg)
                     else:
                         log.verbose(
-                            f"Calling ob {ob.name} intervals update_host for {key}"
+                            f"Calling ob {ob.name} intervals update_host for '{key}'"
                         )
                         ob.intervals.accel_update_host(key)
                 else:
                     log.verbose(
-                        f"Skip update_host for ob {ob.name} intervals {key}, data not present"
+                        f"Skip update_host for ob {ob.name} intervals '{key}', data not present"
                     )
         for key in names["global"]:
             val = self._internal[key]
             if isinstance(val, AcceleratorObject):
                 if not val.accel_in_use():
-                    msg = f"Skipping update_host for {key}, "
-                    msg += "host data in use"
+                    msg = f"Skipping update_host for '{key}', "
+                    msg += "host data already in use"
                     log.verbose(msg)
                 else:
-                    log.verbose(f"Calling Data update_host for {key}")
+                    log.verbose(f"Calling Data update_host for '{key}'")
                     val.accel_update_host()
+            else:
+                msg = f"Data accel_update_host: '{key}' ({type(val)}) "
+                msg += "is not an AcceleratorObject"
+                log.verbose(msg)
 
     def accel_delete(self, names):
         """Delete a specific set of device objects
@@ -625,35 +655,41 @@ def accel_delete(self, names):
         for ob in self.obs:
             for key in names["detdata"]:
                 if ob.detdata.accel_exists(key):
-                    log.verbose(f"Calling ob {ob.name} detdata accel_delete for {key}")
-                    ob.detdata.accel_delete(key)
-                else:
                     log.verbose(
-                        f"Skip delete for ob {ob.name} detdata {key}, data not present"
+                        f"Calling ob {ob.name} detdata accel_delete for '{key}'"
                     )
+                    ob.detdata.accel_delete(key)
+                else:
+                    msg = f"Skip delete for ob {ob.name} detdata '{key}', "
+                    msg += "data not present"
+                    log.verbose(msg)
             for key in names["shared"]:
                 if ob.shared.accel_exists(key):
-                    log.verbose(f"Calling ob {ob.name} shared accel_delete for {key}")
+                    log.verbose(f"Calling ob {ob.name} shared accel_delete for '{key}'")
                     ob.shared.accel_delete(key)
                 else:
-                    log.verbose(
-                        f"Skip delete for ob {ob.name} shared {key}, data not present"
-                    )
+                    msg = f"Skip delete for ob {ob.name} shared '{key}', "
+                    msg += "data not present"
+                    log.verbose(msg)
             for key in names["intervals"]:
                 if ob.intervals.accel_exists(key):
                     log.verbose(
-                        f"Calling ob {ob.name} intervals accel_delete for {key}"
+                        f"Calling ob {ob.name} intervals accel_delete for '{key}'"
                     )
                     ob.intervals.accel_delete(key)
                 else:
-                    log.verbose(
-                        f"Skip delete for ob {ob.name} intervals {key}, data not present"
-                    )
+                    msg = f"Skip delete for ob {ob.name} intervals '{key}', "
+                    msg += "data not present"
+                    log.verbose(msg)
         for key in names["global"]:
             val = self._internal[key]
             if isinstance(val, AcceleratorObject):
-                log.verbose(f"Calling Data accel_delete for {key}")
+                log.verbose(f"Calling Data accel_delete for '{key}'")
                 val.accel_delete()
+            else:
+                msg = f"Data accel_delete: '{key}' ({type(val)}) "
+                msg += "is not an AcceleratorObject"
+                log.verbose(msg)
 
     def accel_clear(self):
         """Delete all accelerator data."""
@@ -665,3 +701,7 @@ def accel_clear(self):
         for key, val in self._internal.items():
             if isinstance(val, AcceleratorObject):
                 val.accel_delete()
+            else:
+                msg = f"Data accel_clear: '{key}' ({type(val)}) "
+                msg += "is not an AcceleratorObject"
+                log.verbose(msg)
diff --git a/src/toast/intervals.py b/src/toast/intervals.py
index 3262db714..d4a33f9f6 100644
--- a/src/toast/intervals.py
+++ b/src/toast/intervals.py
@@ -2,8 +2,7 @@
 # All rights reserved.  Use of this source code is governed by
 # a BSD-style license that can be found in the LICENSE file.
 
-import sys
-from collections.abc import MutableMapping, Sequence
+from collections.abc import Sequence
 
 import numpy as np
 
@@ -14,7 +13,6 @@
     accel_data_present,
     accel_data_update_device,
     accel_data_update_host,
-    accel_enabled,
     use_accel_jax,
     use_accel_omp,
 )
@@ -22,8 +20,7 @@
 from .utils import Logger
 
 if use_accel_jax:
-    import jax
-    import jax.numpy as jnp
+    from .jax.intervals import INTERVALS_JAX
 
 
 def build_interval_dtype():
@@ -119,7 +116,6 @@ def __init__(self, timestamps, intervals=None, timespans=None, samplespans=None)
         else:
             # No data yet
             self.data = np.zeros(0, dtype=interval_dtype).view(np.recarray)
-        self.jax = None
         # This is for the AcceleratorObject mixin class
         self._accel_used = False
 
@@ -359,40 +355,56 @@ def __or__(self, other):
         )
 
     def _accel_exists(self):
+        log = Logger.get()
+        log.verbose(f"IntervalList _accel_exists")
         if len(self.data) == 0:
             return False
+        elif use_accel_omp:
+            return accel_data_present(self.data)
+        elif use_accel_jax:
+            # specialised for the INTERVALS_JAX dtype
+            return isinstance(self.data, INTERVALS_JAX)
         else:
-            if use_accel_omp:
-                return accel_data_present(self.data)
-            elif use_accel_jax:
-                return accel_data_present(self.jax)
-            else:
-                return False
+            return False
 
     def _accel_create(self):
+        log = Logger.get()
+        log.verbose(f"IntervalList _accel_create")
         if use_accel_omp:
             accel_data_create(self.data)
         elif use_accel_jax:
-            accel_data_create(self.jax)
+            # specialised for the INTERVALS_JAX dtype
+            self.data = INTERVALS_JAX(self.data)
 
     def _accel_update_device(self):
+        log = Logger.get()
+        log.verbose(f"IntervalList _accel_update_device")
         if use_accel_omp:
             _ = accel_data_update_device(self.data)
         elif use_accel_jax:
-            self.jax = accel_data_update_device(self.data)
+            # specialised for the INTERVALS_JAX dtype
+            self.data = INTERVALS_JAX(self.data)
 
     def _accel_update_host(self):
+        log = Logger.get()
+        log.verbose(f"IntervalList _accel_update_host")
         if use_accel_omp:
             _ = accel_data_update_host(self.data)
         elif use_accel_jax:
-            self.data = accel_data_update_host(self.jax)
+            # specialised for the INTERVALS_JAX dtype
+            # this moves the data back into a numpy array
+            self.data = self.data.to_host()
 
     def _accel_delete(self):
+        log = Logger.get()
+        log.verbose(f"IntervalList _accel_delete")
         if use_accel_omp:
             accel_data_delete(self.data)
-        elif use_accel_jax:
-            del self.jax
-            self.jax = None
+        elif use_accel_jax and self._accel_exists():
+            # Ensures data has been properly reset
+            # if we observe that its type is still a GPU type
+            # does NOT move data back from GPU
+            self.data = self.data.host_data
 
 
 @function_timer
diff --git a/src/toast/jax/CMakeLists.txt b/src/toast/jax/CMakeLists.txt
new file mode 100644
index 000000000..623e2cf54
--- /dev/null
+++ b/src/toast/jax/CMakeLists.txt
@@ -0,0 +1,12 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    intervals.py
+    mutableArray.py
+    device.py 
+    DESTINATION ${PYTHON_SITE}/toast/jax
+)
+
+add_subdirectory(math)
diff --git a/src/toast/jax/__init__.py b/src/toast/jax/__init__.py
new file mode 100644
index 000000000..c6f667bcd
--- /dev/null
+++ b/src/toast/jax/__init__.py
@@ -0,0 +1,9 @@
+# NOTE:
+# - the code uses [xmap](https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html)
+#   in order to map over named axis for increased readability
+#   however, one could use several vmap to reproduce this functionality
+
+# enable 64bits precision
+from jax.config import config as jax_config
+
+jax_config.update("jax_enable_x64", True)
diff --git a/src/toast/jax/device.py b/src/toast/jax/device.py
new file mode 100644
index 000000000..1ae7dcfa0
--- /dev/null
+++ b/src/toast/jax/device.py
@@ -0,0 +1,56 @@
+import math
+import os
+
+import jax
+
+from .._libtoast import Logger
+
+jax_local_device = None
+"""
+    The device that JAX operators should be using on this process.
+    Use the function `set_JAX_device` to set this value (this should be done during MPI initialization)
+    Use like this in your JAX code: `jax.device_put(data, device=jax_local_device)`
+    Or set globally with: `jax.config.update("jax_default_device", jax_local_device)`
+"""
+
+
+def jax_accel_assign_device(node_procs, node_rank, disabled):
+    """
+    Assign processes to target devices.
+
+    Args:
+        node_procs (int): number of processes per node
+        node_rank (int): rank of the current process, within the node
+        disabled (bool): gpu computing is disabled
+
+    Returns:
+        None: the device is stored in the jax_local_device global variable and jax.config
+    """
+    # list of GPUs / CPUs available
+    devices_available = jax.devices()
+    # gets id of device to be used by this process
+    nb_devices = len(devices_available)
+    device_id = 0 if disabled else (node_rank % nb_devices)
+    # sets the device in a local variable and globaly
+    global jax_local_device
+    jax_local_device = devices_available[device_id]
+    jax.config.update("jax_default_device", jax_local_device)
+    # sets the size of the preallocated memory pool (if it is still possible)
+    default_mem_fraction = 0.99
+    process_per_device = math.ceil(node_procs / nb_devices)
+    mem_fraction = min(default_mem_fraction, default_mem_fraction / process_per_device)
+    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = f"{mem_fraction:.2f}"
+    # displays information on the device picked
+    log = Logger.get()
+    log.debug(
+        f"JAX rank {node_rank}/{node_procs} uses device number {device_id}/{nb_devices} ({jax_local_device})"
+    )
+
+
+def jax_accel_get_device():
+    """Returns the device currenlty used by JAX or errors out."""
+    if jax_local_device is None:
+        raise RuntimeError(
+            "Jax device is not set, please insure that you have called 'accel_assign_device'"
+        )
+    return jax_local_device
diff --git a/src/toast/jax/intervals.py b/src/toast/jax/intervals.py
new file mode 100644
index 000000000..a1c200e15
--- /dev/null
+++ b/src/toast/jax/intervals.py
@@ -0,0 +1,138 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+# TODO we need to clean up the two very similar interval type names
+
+# ------------------------------------------------------------------------------
+# Numpy to JAX intervals
+
+
+class INTERVALS_JAX:
+    """
+    Encapsulate the information from an array of structures into a structure of array
+    each of the inner arrays is converted into a JAX array
+    """
+
+    def __init__(self, data):
+        # otherwise
+        self.size = data.size
+        self.intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(data)
+        self.start = jax.device_put(data.start)
+        self.stop = jax.device_put(data.stop)
+        self.first = jax.device_put(data.first)
+        self.last = jax.device_put(data.last)
+        # sets Numpy buffer aside in case we want to return it
+        self.host_data = data.host_data if isinstance(data, INTERVALS_JAX) else data
+
+    def compute_max_intervals_length(intervals):
+        """
+        given an array of intervals, returns the maximum interval length
+        this function will use a precomputed values when using an INTERVALS_JAX
+        """
+        if isinstance(intervals, INTERVALS_JAX):
+            return intervals.intervals_max_length
+        else:
+            # end+1 as TOAST intervals are inclusive
+            return np.max(1 + intervals.last - intervals.first)
+
+    def to_host(self):
+        """copies data back into the original buffer and returns it"""
+        self.host_data.start[:] = self.start
+        self.host_data.stop[:] = self.stop
+        self.host_data.first[:] = self.first
+        self.host_data.last[:] = self.last
+        return self.host_data
+
+    def __iter__(self):
+        # NOTE: this is only correct if intervals are write only
+        return iter(self.host_data)
+
+    def __len__(self):
+        return self.size
+
+
+# ------------------------------------------------------------------------------
+# Irregular intervals
+
+ALL = slice(None, None, None)
+"""
+Full slice, equivalent to `:` in `[:]`.
+"""
+
+
+class JaxIntervals:
+    """
+    Class to helps dealing with variable-size intervals in JAX.
+    Internally, it pads the data to max_length on read and masks them on write.
+    WARNING: this function is designed to be used inside a jitted-function as it can be very memory hungry otherwise.
+    """
+
+    def __init__(self, starts, ends, max_length):
+        """
+        Builds a JaxIntervals object using the starting and ending points of all intervals
+        plus the length of the larger interval (this needs to be a static quantity).
+        """
+        # 2D tensor of integer of shape (nb_interval,max_length)
+        self.indices = starts[:, jnp.newaxis] + jnp.arange(max_length)
+        # mask that is True on all values that should be ignored in the interval
+        self.mask = self.indices >= ends[:, jnp.newaxis]
+
+    def _interval_of_key(key):
+        """
+        Takes a key, expected to be a JaxIntervals or a tuple with at least one JaxIntervals member
+        and returns (key, mask) where both key and mask can be used on a matrix that would be indexed by the original key.
+        """
+        # insures that the key is a tuple
+        if not isinstance(key, tuple):
+            key = (key,)
+        # insures all elements of the key are valid
+        # and finds the interval
+        mask = None
+
+        def fix_key(key):
+            nonlocal mask
+            if isinstance(key, JaxIntervals):
+                # stores the interval and returns the actual index
+                mask = key.mask
+                return key.indices
+            elif not (mask is None):
+                # adds a trailing dimension to the mask (as there are subsequent dimenssions)
+                mask = jnp.expand_dims(mask, axis=-1)
+            # adds two trailing dimensions to arrays, one per interval dimenssion
+            if (isinstance(key, jnp.ndarray) or isinstance(key, np.ndarray)) and (
+                key.ndim > 0
+            ):
+                return key[:, jnp.newaxis, jnp.newaxis]
+            return key
+
+        key = tuple(fix_key(k) for k in key)
+        # makes sure at least one of the indices was an interval
+        if mask is None:
+            raise RuntimeError(
+                "JaxIntervals: your key should contain a JaxIntervals type."
+            )
+        return (key, mask)
+
+    def get(data, key, padding_value=None):
+        """
+        Gets the data at the given key
+        the result will be padded to keep the interval size constant
+        we will use values from data to pad unless padding_value is not None
+        we expect key to be a JaxIntervals or a tuple with at least one JaxIntervals member.
+        """
+        key, mask = JaxIntervals._interval_of_key(key)
+        return (
+            data[key]
+            if (padding_value is None)
+            else jnp.where(mask, padding_value, data[key])
+        )
+
+    def set(data, key, value_intervals):
+        """
+        Sets the data at the given key with value_intervals
+        we expect key to be a JaxIntervals or a tuple with at least one JaxIntervals member.
+        """
+        key, mask = JaxIntervals._interval_of_key(key)
+        value_interval_masked = jnp.where(mask, data[key], value_intervals)
+        return data.at[key].set(value_interval_masked)
diff --git a/src/toast/jax/math/CMakeLists.txt b/src/toast/jax/math/CMakeLists.txt
new file mode 100644
index 000000000..42dca37bf
--- /dev/null
+++ b/src/toast/jax/math/CMakeLists.txt
@@ -0,0 +1,10 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    qarray.py
+    healpix.py
+    DESTINATION ${PYTHON_SITE}/toast/jax/math
+)
+
diff --git a/src/toast/jax/math/__init__.py b/src/toast/jax/math/__init__.py
new file mode 100644
index 000000000..3f3c118b0
--- /dev/null
+++ b/src/toast/jax/math/__init__.py
@@ -0,0 +1,3 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
diff --git a/src/toast/jax/math/healpix.py b/src/toast/jax/math/healpix.py
new file mode 100644
index 000000000..e6de6d77d
--- /dev/null
+++ b/src/toast/jax/math/healpix.py
@@ -0,0 +1,236 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+TWOINVPI = 0.63661977236758134308
+MACHINE_EPSILON = np.finfo(np.float64).eps
+
+
+class HPIX_JAX:
+    """
+    JAX compatible HPIX structure
+    This class has an efficient hash that lets it be cached
+    """
+
+    def __init__(self, nside):
+        self.nside = nside
+        self.ncap = 2 * (nside * nside - nside)
+        self.npix = 12 * nside * nside
+        self.dnside = float(nside)
+        self.twonside = 2 * nside
+        self.fournside = 4 * nside
+        self.nsideplusone = nside + 1
+        self.nsideminusone = nside - 1
+        self.halfnside = 0.5 * self.dnside
+        self.tqnside = 0.75 * self.dnside
+
+        factor = 0
+        while nside != (1 << factor):
+            factor += 1
+        self.factor = factor
+
+        self.jr = jnp.array([2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4])
+        self.jq = jnp.array([1, 3, 5, 7, 0, 2, 4, 6, 1, 3, 5, 7])
+
+        m = jnp.arange(start=0, stop=0x100)
+        self.utab = (
+            (m & 0x1)
+            | ((m & 0x2) << 1)
+            | ((m & 0x4) << 2)
+            | ((m & 0x8) << 3)
+            | ((m & 0x10) << 4)
+            | ((m & 0x20) << 5)
+            | ((m & 0x40) << 6)
+            | ((m & 0x80) << 7)
+        )
+        self.ctab = (
+            (m & 0x1)
+            | ((m & 0x2) << 7)
+            | ((m & 0x4) >> 1)
+            | ((m & 0x8) << 6)
+            | ((m & 0x10) >> 2)
+            | ((m & 0x20) << 5)
+            | ((m & 0x40) >> 3)
+            | ((m & 0x80) << 4)
+        )
+
+    def xy2pix(self, x, y):
+        return (
+            self.utab[x & 0xFF]
+            | (self.utab[(x >> 8) & 0xFF] << 16)
+            | (self.utab[(x >> 16) & 0xFF] << 32)
+            | (self.utab[(x >> 24) & 0xFF] << 48)
+            | (self.utab[y & 0xFF] << 1)
+            | (self.utab[(y >> 8) & 0xFF] << 17)
+            | (self.utab[(y >> 16) & 0xFF] << 33)
+            | (self.utab[(y >> 24) & 0xFF] << 49)
+        )
+
+    def __key(self):
+        return self.nside
+
+    def __hash__(self):
+        return hash(self.__key())
+
+    def __eq__(self, other):
+        if isinstance(other, HPIX_JAX):
+            return self.__key() == other.__key()
+        return NotImplemented
+
+
+def zphi2nest(hpix, phi, region, z, rtz):
+    """
+    Args:
+        hpix (HealpixPixels):  The healpix projection object.
+        phi (double)
+        region (int)
+        z (double)
+        rtz (double)
+
+    Returns:
+        pix (int)
+    """
+    phi = jnp.where(jnp.abs(phi) < MACHINE_EPSILON, 0.0, phi)
+
+    tt = phi * TWOINVPI
+    tt = jnp.where(phi < 0.0, tt + 4.0, tt)
+
+    # NOTE: this is very slightly faster than a jnp.where here
+    # then branch
+    def then_branch(tt, rtz, z):
+        temp1 = hpix.halfnside + hpix.dnside * tt
+        temp2 = hpix.tqnside * z
+
+        jp = jnp.int64(temp1 - temp2)
+        jm = jnp.int64(temp1 + temp2)
+
+        ifp = jp >> hpix.factor
+        ifm = jm >> hpix.factor
+
+        face = jnp.where(
+            ifp == ifm,
+            jnp.where(ifp == 4, 4, ifp + 4),
+            jnp.where(ifp < ifm, ifp, ifm + 8),
+        )
+
+        x = jm & hpix.nsideminusone
+        y = hpix.nsideminusone - (jp & hpix.nsideminusone)
+        return (face, x, y)
+
+    # else branch
+    def else_branch(tt, rtz, z):
+        ntt = jnp.int64(tt)
+
+        tp = tt - jnp.double(ntt)
+
+        temp1 = hpix.dnside * rtz
+
+        jp = jnp.int64(tp * temp1)
+        jp = jnp.where(jp >= hpix.nside, hpix.nsideminusone, jp)
+
+        jm = jnp.int64((1.0 - tp) * temp1)
+        jm = jnp.where(jm >= hpix.nside, hpix.nsideminusone, jm)
+
+        face = jnp.where(z >= 0, ntt, ntt + 8)
+        x = jnp.where(z >= 0, hpix.nsideminusone - jm, jp)
+        y = jnp.where(z >= 0, hpix.nsideminusone - jp, jm)
+        return (face, x, y)
+
+    # test
+    (face, x, y) = jax.lax.cond(
+        jnp.abs(region) == 1, then_branch, else_branch, tt, rtz, z
+    )
+
+    sipf = hpix.xy2pix(x, y)
+    pix = jnp.int64(sipf) + (face << (2 * hpix.factor))
+    return pix
+
+
+def zphi2ring(hpix, phi, region, z, rtz):
+    """
+    Args:
+        hpix (HealpixPixels):  The healpix projection object.
+        phi (double)
+        region (int)
+        z (double)
+        rtz (double)
+
+    Returns:
+        pix (int)
+    """
+    phi = jnp.where(jnp.abs(phi) < MACHINE_EPSILON, 0.0, phi)
+
+    tt = phi * TWOINVPI
+    tt = jnp.where(phi < 0.0, tt + 4.0, tt)
+
+    # NOTE: this is very slightly faster than a jnp.where here
+    # then branch
+    def then_branch(tt, rtz, z):
+        temp1 = hpix.halfnside + hpix.dnside * tt
+        temp2 = hpix.tqnside * z
+
+        jp = jnp.int64(temp1 - temp2)
+        jm = jnp.int64(temp1 + temp2)
+
+        ir = hpix.nsideplusone + jp - jm
+        kshift = 1 - (ir & 1)
+
+        ip = (jp + jm - hpix.nside + kshift + 1) >> 1
+        ip = ip % hpix.fournside
+
+        pix = hpix.ncap + ((ir - 1) * hpix.fournside + ip)
+        return pix
+
+    # else branch
+    def else_branch(tt, rtz, z):
+        tp = tt - jnp.floor(tt)
+
+        temp1 = hpix.dnside * rtz
+
+        jp = jnp.int64(tp * temp1)
+        jm = jnp.int64((1.0 - tp) * temp1)
+        ir = jp + jm + 1
+        ip = jnp.int64(tt * jnp.double(ir))
+        longpart = jnp.int64(ip / (4 * ir))
+        ip = ip - longpart
+
+        pix_pos = 2 * ir * (ir - 1) + ip
+        pix_neg = hpix.npix - 2 * ir * (ir + 1) + ip
+        pix = jnp.where(region > 0, pix_pos, pix_neg)
+        return pix
+
+    # test
+    pix = jax.lax.cond(jnp.abs(region) == 1, then_branch, else_branch, tt, rtz, z)
+
+    return pix
+
+
+def vec2zphi(vec):
+    """
+    Args:
+        vec (array, double) of size 3
+
+    Returns:
+        (phi, region, z, rtz)
+        phi (double)
+        region (int)
+        z (double)
+        rtz (double)
+    """
+    z = vec[2]
+    za = jnp.abs(z)
+
+    # region encodes BOTH the sign of Z and whether its
+    # absolute value is greater than 2/3.
+    itemps = jnp.where(z > 0.0, 1, -1)
+    region = jnp.where(za <= 2.0 / 3.0, itemps, 2 * itemps)
+
+    tz = 3.0 * (1.0 - za)
+    rtz = jnp.sqrt(tz)
+    phi = jnp.arctan2(vec[1], vec[0])
+
+    return (phi, region, z, rtz)
diff --git a/src/toast/jax/math/qarray.py b/src/toast/jax/math/qarray.py
new file mode 100644
index 000000000..169e6e6a2
--- /dev/null
+++ b/src/toast/jax/math/qarray.py
@@ -0,0 +1,132 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..mutableArray import MutableJaxArray
+
+
+def rotate_one_one(q, v_in):
+    """
+    Rotate a vector by a quaternion.
+
+    Args:
+        q(array, double): quaternion of shape (4)
+        v_in(array, double): vector of size 3
+
+    Returns:
+        v_out(array, double): vector of size 3
+    """
+    # normalize quaternion
+    q_unit = q / jnp.linalg.norm(q)
+
+    # builds the elements that make the matrix representation of the quaternion
+    xw = q_unit[3] * q_unit[0]
+    yw = q_unit[3] * q_unit[1]
+    zw = q_unit[3] * q_unit[2]
+    x2 = -q_unit[0] * q_unit[0]
+    xy = q_unit[0] * q_unit[1]
+    xz = q_unit[0] * q_unit[2]
+    y2 = -q_unit[1] * q_unit[1]
+    yz = q_unit[1] * q_unit[2]
+    z2 = -q_unit[2] * q_unit[2]
+
+    # matrix product
+    v_out_0 = (
+        2 * ((y2 + z2) * v_in[0] + (xy - zw) * v_in[1] + (yw + xz) * v_in[2]) + v_in[0]
+    )
+    v_out_1 = (
+        2 * ((zw + xy) * v_in[0] + (x2 + z2) * v_in[1] + (yz - xw) * v_in[2]) + v_in[1]
+    )
+    v_out_2 = (
+        2 * ((xz - yw) * v_in[0] + (xw + yz) * v_in[1] + (x2 + y2) * v_in[2]) + v_in[2]
+    )
+
+    return jnp.array([v_out_0, v_out_1, v_out_2])
+
+
+# def mult_one_one(p, q):
+#    """
+#    compose two quaternions
+#
+#    NOTE: this version uses an explicit shuffling of the coordinates
+#    """
+#    r0 =  p[0] * q[3] + p[1] * q[2] - p[2] * q[1] + p[3] * q[0]
+#    r1 = -p[0] * q[2] + p[1] * q[3] + p[2] * q[0] + p[3] * q[1]
+#    r2 =  p[0] * q[1] - p[1] * q[0] + p[2] * q[3] + p[3] * q[2]
+#    r3 = -p[0] * q[0] - p[1] * q[1] - p[2] * q[2] + p[3] * q[3]
+#    return jnp.array([r0,r1,r2,r3])
+
+
+def mult_one_one(p, q):
+    """
+    compose two quaternions
+
+    NOTE: this version use a tensor product to keep computation to matrix products
+    """
+    # reshuffles q into a matrix
+    mat = jnp.array(
+        [
+            [[0, 0, 0, -1], [0, 0, 1, 0], [0, -1, 0, 0], [1, 0, 0, 0]],  # row1
+            [[0, 0, -1, 0], [0, 0, 0, -1], [1, 0, 0, 0], [0, 1, 0, 0]],  # row2
+            [[0, 1, 0, 0], [-1, 0, 0, 0], [0, 0, 0, -1], [0, 0, 1, 0]],  # row3
+            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],  # row4
+        ]
+    )
+    qMat = jnp.dot(q, mat)
+    # multiplies p with the reshuffled q
+    return jnp.dot(p, qMat)
+
+
+# loops on either or both of the inputs
+mult_one_many = jax.vmap(mult_one_one, in_axes=(None, 0), out_axes=0)
+mult_many_one = jax.vmap(mult_one_one, in_axes=(0, None), out_axes=0)
+mult_many_many = jax.vmap(mult_one_one, in_axes=(0, 0), out_axes=0)
+
+
+def mult_pure(p_in, q_in):
+    # picks the correct impelmentation depending on which input is an array (if any)
+    p_is_array = p_in.ndim > 1
+    q_is_array = q_in.ndim > 1
+    if p_is_array and q_is_array:
+        return mult_many_many(p_in, q_in)
+    if p_is_array:
+        return mult_many_one(p_in, q_in)
+    if q_is_array:
+        return mult_one_many(p_in, q_in)
+    return mult_one_one(p_in, q_in)
+
+
+# jit compiles the jax function
+mult_pure = jax.jit(mult_pure)
+
+
+def mult(p_in, q_in):
+    """
+    Multiply quaternion arrays.
+
+    The number of quaternions in both input arrays should either be
+    equal or there should be one quaternion in one of the arrays (which
+    is then multiplied by all quaternions in the other array).
+
+    Args:
+        p_in (array_like):  flattened 1D array of float64 values.
+        q_in (array_like):  flattened 1D array of float64 values.
+
+    Returns:
+        out (array_like):  flattened 1D array of float64 values.
+    """
+    # casts the data to arrays and performs the computation
+    p_in_input = MutableJaxArray.to_array(p_in)
+    q_in_input = MutableJaxArray.to_array(q_in)
+    out = mult_pure(p_in_input, q_in_input)
+    # converts to a numpy type if the input was a numpy type
+    return jax.device_get(out) if isinstance(p_in, np.ndarray) else out
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("qarray");'
diff --git a/src/toast/jax/mutableArray.py b/src/toast/jax/mutableArray.py
new file mode 100644
index 000000000..0873a7ca8
--- /dev/null
+++ b/src/toast/jax/mutableArray.py
@@ -0,0 +1,113 @@
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from pshmem import MPIShared
+
+from ..utils import AlignedF64, AlignedI64
+
+
+class MutableJaxArray:
+    """
+    This class encapsulate a jax array to give the illusion of mutability
+    simplifying integration within toast
+    It is NOT designed for computation but, rather, as a container
+    """
+
+    data: jnp.DeviceArray
+    shape: Tuple
+    dtype: np.dtype
+    nbytes: np.int64
+
+    def __init__(self, cpu_data, gpu_data=None):
+        """
+        encapsulate an array as a jax array
+        you can pass `gpu_data` (properly shaped) if you want to avoid a data transfert
+        """
+        # stores cpu and gpu data
+        if isinstance(cpu_data, MutableJaxArray):
+            self.host_data = cpu_data.host_data
+            self.data = cpu_data.data
+        else:
+            self.host_data = cpu_data
+            self.data = gpu_data
+            if gpu_data is None:
+                data = MutableJaxArray.to_array(cpu_data)
+                self.data = jax.device_put(data)
+
+        # gets basic information on the data
+        self.shape = self.data.shape
+        self.size = self.data.size
+        self.dtype = self.data.dtype
+        self.nbytes = self.data.nbytes
+
+    def to_array(input):
+        """
+        Casts all array types used in TOAST to JAX-compatible array types
+        (either a numpy array or a jax array)
+        """
+        if isinstance(input, MPIShared):
+            # this type needs to be cast to its inner type
+            # NOTE: the inner type might be a mutable array or a normal array
+            input = input.data
+        if isinstance(input, MutableJaxArray):
+            # those types need to be cast into their inner data
+            return input.data
+        elif isinstance(input, AlignedI64) or isinstance(input, AlignedF64):
+            # NOTE: get inner array field, raw numpy conversion is very expensive
+            return input.array()
+        elif isinstance(input, np.ndarray) or isinstance(input, jax.numpy.ndarray):
+            # those types can be feed to JAX raw with no error or performance problems
+            return input
+        else:
+            # errors-out on other datatypes
+            # so that we can make sure we are using the most efficient convertion available
+            raise RuntimeError(
+                f"Passed a {type(input)} to MutableJaxArray.to_array. Please find the best way to convert it to a Numpy array and update the function."
+            )
+
+    def to_host(self):
+        """
+        updates the original host container and returns it
+        NOTE: we purposefully do not overload __array__ to avoid accidental conversions
+        """
+        self.host_data[:] = self.data
+        return self.host_data
+
+    def __setitem__(self, key, value):
+        """
+        updates the inner array in place
+        """
+        if key == slice(None):
+            # optimises the [:] case
+            self.data = value
+        else:
+            self.data = self.data.at[key].set(value)
+
+    def __getitem__(self, key):
+        """
+        access the inner array
+        """
+        return self.data[key]
+
+    def reshape(self, shape):
+        """
+        produces a new mutable array with a different shape
+        the array has a copy of the gpu data (changes will not be propagated to the original)
+        and a reshape of the cpu data (change might be propagated to the original, depending on numpy's behaviour)
+        """
+        reshaped_cpu_data = np.reshape(self.host_data, newshape=shape)
+        reshaped_gpu_data = jnp.reshape(self.data, newshape=shape)
+        return MutableJaxArray(reshaped_cpu_data, reshaped_gpu_data)
+
+    def __str__(self):
+        """
+        returns a string representation of the content of the array
+        """
+        return self.data.__str__()
+
+    def __eq__(self, other):
+        raise RuntimeError(
+            "MutableJaxArray: tried an equality test on a MutableJaxArray. This container is not designed for computations, you likely have a data movement bug somewhere in your program."
+        )
diff --git a/src/toast/mpi.py b/src/toast/mpi.py
index 31a45e745..bdb336b4a 100644
--- a/src/toast/mpi.py
+++ b/src/toast/mpi.py
@@ -5,7 +5,7 @@
 import os
 import time
 
-from ._libtoast import Environment, Logger
+from ._libtoast import Logger
 
 use_mpi = None
 MPI = None
@@ -43,10 +43,8 @@ class MPI_Comm:
                 log.debug("mpi4py not found- using serial operations only")
                 use_mpi = False
 
-from ._libtoast import accel_assign_device
-
 # Assign each process to an accelerator device
-from .accelerator import jax_local_device, use_accel_jax, use_accel_omp
+from .accelerator import accel_assign_device, use_accel_jax, use_accel_omp
 
 if use_accel_omp or use_accel_jax:
     node_procs = 1
@@ -59,20 +57,8 @@ class MPI_Comm:
         accel_assign_device(node_procs, node_rank, False)
         nodecomm.Free()
         del nodecomm
-    if use_accel_omp:
-        accel_assign_device(node_procs, node_rank, False)
     else:
-        import jax
-
-        n_target = len(jax.devices())
-        if n_target == 0:
-            jax_local_device = 0
-        else:
-            proc_per_dev = node_procs // n_target
-            if n_target * proc_per_dev < node_procs:
-                proc_per_dev += 1
-            target_dev = node_rank // proc_per_dev
-            jax_local_device = jax.devices()[target_dev]
+        accel_assign_device(node_procs, node_rank, False)
 else:
     # Disabled == True
     accel_assign_device(1, 0, True)
diff --git a/src/toast/observation_data.py b/src/toast/observation_data.py
index 2097a31a5..8a5dce8ed 100644
--- a/src/toast/observation_data.py
+++ b/src/toast/observation_data.py
@@ -2,7 +2,6 @@
 # All rights reserved.  Use of this source code is governed by
 # a BSD-style license that can be found in the LICENSE file.
 
-import sys
 from collections.abc import Mapping, MutableMapping
 from typing import NamedTuple
 
@@ -22,26 +21,13 @@
     use_accel_omp,
 )
 from .intervals import IntervalList
-from .mpi import MPI, comm_equal, comm_equivalent
-from .timing import function_timer
-from .utils import (
-    AlignedF32,
-    AlignedF64,
-    AlignedI8,
-    AlignedI16,
-    AlignedI32,
-    AlignedI64,
-    AlignedU8,
-    AlignedU16,
-    AlignedU32,
-    AlignedU64,
-    Logger,
-    dtype_to_aligned,
-)
+from .mpi import MPI, comm_equivalent
+from .utils import Logger, dtype_to_aligned
 
 if use_accel_jax:
     import jax
-    import jax.numpy as jnp
+
+    from .jax.mutableArray import MutableJaxArray
 
 
 class DetectorData(AcceleratorObject):
@@ -106,7 +92,6 @@ def __init__(
         self._fullsize = 0
         self._memsize = 0
         self._raw = None
-        self._raw_jax = None
 
         if view_data is None:
             # Allocate the data
@@ -271,12 +256,21 @@ def change_detectors(self, detectors):
             # We can re-use the existing memory
             self._shape = shp
             self._flatshape = flatshape
+            # reinitialise _data
+            should_update_GPU = self.accel_exists()
             self._flatdata = self._raw.array()[: self._flatshape]
             self._flatdata[:] = 0
             self._data = self._flatdata.reshape(self._shape)
-            if self.accel_exists():
-                # Should we zero the device memory?
-                pass
+            # move things to GPU if needed
+            if should_update_GPU:
+                if use_accel_jax:
+                    # NOTE: we create an array of zeroes directly on GPU to avoid useless data movement
+                    self._data = MutableJaxArray(
+                        self._data, gpu_data=jax.numpy.zeros_like(self._data)
+                    )
+                else:
+                    # FIXME: Should we zero the device memory?
+                    pass
             realloced = False
         return realloced
 
@@ -290,6 +284,7 @@ def clear(self):
         """
         if hasattr(self, "_data"):
             del self._data
+            self._data = None
         if hasattr(self, "_is_view") and not self._is_view:
             if hasattr(self, "_flatdata"):
                 del self._flatdata
@@ -303,9 +298,6 @@ def clear(self):
                     # since that function also manipulates self._raw.
                     if use_accel_omp:
                         accel_data_delete(self._raw)
-                    elif use_accel_jax:
-                        del self._raw_jax
-                        self._raw_jax = None
                 if self._raw is not None:
                     self._raw.clear()
                 del self._raw
@@ -453,6 +445,8 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
     def _accel_exists(self):
+        log = Logger.get()
+        log.verbose(f"DetectorData _accel_exists")
         if self._raw is None:
             # We have a view
             return False
@@ -460,40 +454,45 @@ def _accel_exists(self):
             if use_accel_omp:
                 return accel_data_present(self._raw)
             elif use_accel_jax:
-                return accel_data_present(self._raw_jax)
+                return accel_data_present(self._data)
             else:
                 return False
 
     def _accel_create(self):
+        log = Logger.get()
+        log.verbose(f"DetectorData _accel_create")
         if use_accel_omp:
             accel_data_create(self._raw)
         elif use_accel_jax:
-            accel_data_create(self._raw_jax)
+            self._data = accel_data_create(self._data)
 
     def _accel_update_device(self):
+        log = Logger.get()
+        log.verbose(f"DetectorData _accel_update_device")
         if use_accel_omp:
             _ = accel_data_update_device(self._raw)
         elif use_accel_jax:
-            self._raw_jax = accel_data_update_device(self._raw.array())
-            self._flatdata = self._raw_jax[: self._flatshape]
-            self._data = self._flatdata.reshape(self._shape)
+            self._data = accel_data_update_device(self._data)
 
     def _accel_update_host(self):
+        log = Logger.get()
+        log.verbose(f"DetectorData _accel_update_host")
         if use_accel_omp:
             _ = accel_data_update_host(self._raw)
         elif use_accel_jax:
-            self._raw[:] = accel_data_update_host(self._raw_jax)
-            self._flatdata = self._raw.array()[: self._flatshape]
-            self._data = self._flatdata.reshape(self._shape)
+            self._data = accel_data_update_host(self._data)
 
     def _accel_delete(self):
+        log = Logger.get()
+        log.verbose(f"DetectorData _accel_delete")
         if use_accel_omp:
             accel_data_delete(self._raw)
-        elif use_accel_jax:
-            del self._raw_jax
-            self._raw_jax = None
-            self._flatdata = self._raw.array()[: self._flatshape]
-            self._data = self._flatdata.reshape(self._shape)
+        elif use_accel_jax and self._accel_exists():
+            # Ensures _data has been properly reset
+            # if we observe that its types is still a GPU types
+            # does NOT move data back from GPU
+            # using _raw should be equivalent
+            self._data = self._data.host_data
 
 
 class DetDataManager(MutableMapping):
@@ -752,7 +751,9 @@ def accel_exists(self, key):
             return False
         log = Logger.get()
         result = self._internal[key].accel_exists()
-        log.verbose(f"DetDataMgr {key} accel_exists = {result}")
+        msg = f"DetDataMgr {key} type = {type(self._internal[key])} "
+        msg += f"accel_exists = {result}"
+        log.verbose(msg)
         return result
 
     def accel_in_use(self, key):
@@ -796,7 +797,7 @@ def accel_create(self, key):
         if not accel_enabled():
             return
         log = Logger.get()
-        log.verbose(f"DetDataMgr {key} accel_create")
+        log.verbose(f"DetDataMgr {key} type = {type(self._internal[key])} accel_create")
         self._internal[key].accel_create()
 
     def accel_update_device(self, key):
@@ -812,7 +813,9 @@ def accel_update_device(self, key):
         if not accel_enabled():
             return
         log = Logger.get()
-        log.verbose(f"DetDataMgr {key} accel_update_device")
+        log.verbose(
+            f"DetDataMgr {key} type = {type(self._internal[key])} accel_update_device"
+        )
         self._internal[key].accel_update_device()
 
     def accel_update_host(self, key):
@@ -828,7 +831,9 @@ def accel_update_host(self, key):
         if not accel_enabled():
             return
         log = Logger.get()
-        log.verbose(f"DetDataMgr {key} accel_update_host")
+        log.verbose(
+            f"DetDataMgr {key} type = {type(self._internal[key])} accel_update_host"
+        )
         self._internal[key].accel_update_host()
 
     def accel_delete(self, key):
@@ -845,10 +850,11 @@ def accel_delete(self, key):
         if not accel_enabled():
             return
         if not self._internal[key].accel_exists():
-            msg = f"Detector data '{key}' is not present on device, cannot delete"
+            msg = f"Detector data '{key}' type = {type(self._internal[key])} "
+            msg += f"is not present on device, cannot delete"
             log.error(msg)
             raise RuntimeError(msg)
-        log.verbose(f"DetDataMgr {key} accel_delete")
+        log.verbose(f"DetDataMgr {key} type = {type(self._internal[key])} accel_delete")
         self._internal[key].accel_delete()
 
     def accel_clear(self):
@@ -863,7 +869,9 @@ def accel_clear(self):
         log = Logger.get()
         for key in self._internal:
             if self._internal[key].accel_exists():
-                log.verbose(f"DetDataMgr {key} accel_delete")
+                log.verbose(
+                    f"DetDataMgr {key} type = {type(self._internal[key])} accel_delete"
+                )
                 self._internal[key].accel_delete()
 
     # Mapping methods
@@ -1151,7 +1159,6 @@ def __init__(self, dist):
         # data object and a string specifying which communicator it
         # is distributed over:  "group", "row", or "column".
         self._internal = dict()
-        self.jax = dict()
         self._accel_used = dict()
 
     def create_group(self, name, shape, dtype=None):
@@ -1389,7 +1396,7 @@ def accel_exists(self, key):
         if use_accel_omp:
             result = accel_data_present(self._internal[key].shdata._flat)
         elif use_accel_jax:
-            result = key in self.jax
+            result = accel_data_present(self._internal[key].shdata.data)
         else:
             result = False
 
@@ -1455,6 +1462,10 @@ def accel_create(self, key):
         log.verbose(f"SharedDataMgr {key} accel_create")
         if use_accel_omp:
             accel_data_create(self._internal[key].shdata._flat)
+        elif use_accel_jax:
+            self._internal[key].shdata.data = MutableJaxArray(
+                self._internal[key].shdata.data
+            )
 
     def accel_update_device(self, key):
         """Copy the named shared data to the accelerator.
@@ -1488,8 +1499,9 @@ def accel_update_device(self, key):
         if use_accel_omp:
             _ = accel_data_update_device(self._internal[key].shdata._flat)
         elif use_accel_jax:
-            # FIXME: add device ID here
-            self.jax[key] = jax.device_put(self._internal[key].shdata._flat)
+            self._internal[key].shdata.data = MutableJaxArray(
+                self._internal[key].shdata.data
+            )
 
         self._accel_used[key] = True
 
@@ -1525,13 +1537,9 @@ def accel_update_host(self, key):
         if use_accel_omp:
             _ = accel_data_update_host(self._internal[key].shdata._flat)
         elif use_accel_jax:
-            rnk = 0
-            if self._internal[key].shdata.comm is not None:
-                rnk = self._internal[key].shdata.comm.rank
-            dt = None
-            if rnk == 0:
-                dt = self.jax[key].copy()
-            self._internal[key].shdata.set(dt, fromrank=0)
+            self._internal[key].shdata.data = accel_data_update_host(
+                self._internal[key].shdata.data
+            )
 
         self._accel_used[key] = False
 
@@ -1561,7 +1569,9 @@ def accel_delete(self, key):
         if use_accel_omp:
             accel_data_delete(self._internal[key].shdata._flat)
         elif use_accel_jax:
-            del self.jax[key]
+            self._internal[key].shdata.data = accel_data_delete(
+                self._internal[key].shdata.data
+            )
 
         self._accel_used[key] = False
 
@@ -1581,7 +1591,9 @@ def accel_clear(self):
                 if use_accel_omp:
                     accel_data_delete(self._internal[key].shdata._flat)
                 elif use_accel_jax:
-                    del self.jax[key]
+                    self._internal[key].shdata.data = accel_data_delete(
+                        self._internal[key].shdata.data
+                    )
             self._accel_used[key] = False
 
     # Mapping methods
@@ -1656,9 +1668,9 @@ def _assign_array(self, key, value, commtype):
         off = None
         if myrank == fromrank:
             if value.shape != self._internal[key].shdata.shape:
-                raise ValueError(
-                    "When assigning directly to a shared object, the value must have the same dimensions"
-                )
+                msg = "When assigning directly to a shared object, the value "
+                msg += "must have the same dimensions"
+                raise ValueError(msg)
             off = tuple([0 for x in self._internal[key].shdata.shape])
         self._internal[key].shdata.set(value, offset=off, fromrank=fromrank)
 
diff --git a/src/toast/ops/CMakeLists.txt b/src/toast/ops/CMakeLists.txt
index f4e290421..2900629cb 100644
--- a/src/toast/ops/CMakeLists.txt
+++ b/src/toast/ops/CMakeLists.txt
@@ -23,14 +23,10 @@ install(FILES
     flag_intervals.py
     common_mode_noise.py
     noise_model.py
-    noise_weight.py
     gainscrambler.py
     pointing.py
-    pointing_detector.py
-    scan_map.py
     scan_healpix.py
     scan_wcs.py
-    mapmaker_utils.py
     mapmaker_binning.py
     mapmaker_solve.py
     mapmaker_templates.py
@@ -41,7 +37,6 @@ install(FILES
     sim_gaindrifts.py
     sim_crosstalk.py
     totalconvolve.py
-    polyfilter.py
     groundfilter.py
     hwpfilter.py
     statistics.py
@@ -55,8 +50,6 @@ install(FILES
     crosslinking.py
     sss.py
     demodulation.py
-    stokes_weights.py
-    pixels_healpix.py
     pixels_wcs.py
     filterbin.py
     save_hdf5.py
@@ -66,3 +59,11 @@ install(FILES
     yield_cut.py
     DESTINATION ${PYTHON_SITE}/toast/ops
 )
+
+add_subdirectory(pointing_detector)
+add_subdirectory(pixels_healpix)
+add_subdirectory(stokes_weights)
+add_subdirectory(mapmaker_utils)
+add_subdirectory(scan_map)
+add_subdirectory(noise_weight)
+add_subdirectory(polyfilter)
diff --git a/src/toast/ops/conviqt.py b/src/toast/ops/conviqt.py
index 6e6a5c424..0df8a6cf1 100644
--- a/src/toast/ops/conviqt.py
+++ b/src/toast/ops/conviqt.py
@@ -261,9 +261,12 @@ def _exec(self, data, detectors=None, **kwargs):
         timer = Timer()
         timer.start()
 
-        all_detectors = self._get_all_detectors(data, detectors)
-
         self.units = data.detector_units(self.det_data)
+        if self.units is None:
+            # This means that the data does not yet exist
+            self.units = self.det_data_units
+
+        all_detectors = self._get_all_detectors(data, detectors)
 
         for det in all_detectors:
             verbose = self.comm.rank == 0 and self.verbosity > 0
@@ -315,7 +318,7 @@ def _get_all_detectors(self, data, detectors):
                 my_dets.add(det)
             # Make sure detector data output exists
             exists = obs.detdata.ensure(
-                self.det_data, detectors=detectors, create_units=self.det_data_units
+                self.det_data, detectors=detectors, create_units=self.units
             )
         all_dets = self.comm.gather(my_dets, root=0)
         if self.comm.rank == 0:
@@ -554,7 +557,7 @@ def calibrate_signal(self, data, det, beam, convolved_data, verbose):
             epsilon = self._get_epsilon(focalplane, det)
             # Make sure detector data output exists
             exists = obs.detdata.ensure(
-                self.det_data, detectors=[det], create_units=self.det_data_units
+                self.det_data, detectors=[det], create_units=self.units
             )
             # Loop over views
             views = obs.view[self.view]
@@ -571,7 +574,7 @@ def save(self, data, det, convolved_data, verbose):
         timer = Timer()
         timer.start()
         offset = 0
-        scale = unit_conversion(u.K, self.det_data_units)
+        scale = unit_conversion(u.K, self.units)
         for obs in data.obs:
             if det not in obs.local_detectors:
                 continue
@@ -630,6 +633,11 @@ def _exec(self, data, detectors=None, **kwargs):
         timer = Timer()
         timer.start()
 
+        self.units = data.detector_units(self.det_data)
+        if self.units is None:
+            # This means that the data does not yet exist
+            self.units = self.det_data_units
+
         # Expand detector pointing
         self.detector_pointing.apply(data, detectors=detectors)
 
@@ -745,6 +753,11 @@ def _exec(self, data, detectors=None, **kwargs):
         timer = Timer()
         timer.start()
 
+        self.units = data.detector_units(self.det_data)
+        if self.units is None:
+            # This means that the data does not yet exist
+            self.units = self.det_data_units
+
         # Expand detector pointing
         self.detector_pointing.apply(data, detectors=detectors)
 
diff --git a/src/toast/ops/mapmaker_templates.py b/src/toast/ops/mapmaker_templates.py
index 7d6a9d7ff..f6e7195ff 100644
--- a/src/toast/ops/mapmaker_templates.py
+++ b/src/toast/ops/mapmaker_templates.py
@@ -3,12 +3,12 @@
 # a BSD-style license that can be found in the LICENSE file.
 
 import os
-from collections import OrderedDict
 
 import numpy as np
 import traitlets
 from astropy import units as u
 
+from ..accelerator import ImplementationType
 from ..mpi import MPI
 from ..observation import default_values as defaults
 from ..pixels import PixelData
@@ -20,13 +20,12 @@
 from ..utils import Logger
 from .arithmetic import Combine
 from .copy import Copy
-from .delete import Delete
 from .mapmaker_solve import SolverLHS, SolverRHS, solve
 from .mapmaker_utils import CovarianceAndHits
 from .memory_counter import MemoryCounter
 from .operator import Operator
 from .pipeline import Pipeline
-from .scan_map import ScanMap, ScanMask
+from .scan_map import ScanMask
 
 
 @trait_docs
@@ -107,7 +106,7 @@ def duplicate(self):
         ret._initialized = self._initialized
         return ret
 
-    def apply_precond(self, amps_in, amps_out):
+    def apply_precond(self, amps_in, amps_out, use_accel=False, **kwargs):
         """Apply the preconditioner from all templates to the amplitudes.
 
         This can only be called after the operator has been used at least once so that
@@ -126,9 +125,9 @@ def apply_precond(self, amps_in, amps_out):
                 "You must call exec() once before applying preconditioners"
             )
         for tmpl in self.templates:
-            tmpl.apply_precond(amps_in[tmpl.name], amps_out[tmpl.name])
+            tmpl.apply_precond(amps_in[tmpl.name], amps_out[tmpl.name], use_accel=use_accel, **kwargs)
 
-    def add_prior(self, amps_in, amps_out):
+    def add_prior(self, amps_in, amps_out, use_accel=False, **kwargs):
         """Apply the noise prior from all templates to the amplitudes.
 
         This can only be called after the operator has been used at least once so that
@@ -147,7 +146,7 @@ def add_prior(self, amps_in, amps_out):
                 "You must call exec() once before applying the noise prior"
             )
         for tmpl in self.templates:
-            tmpl.add_prior(amps_in[tmpl.name], amps_out[tmpl.name])
+            tmpl.add_prior(amps_in[tmpl.name], amps_out[tmpl.name], use_accel=use_accel, **kwargs)
 
     @property
     def n_enabled_templates(self):
@@ -165,6 +164,9 @@ def reset_templates(self):
     def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         # Check that the detector data is set
         if self.det_data is None:
             raise RuntimeError("You must set the det_data trait before calling exec()")
@@ -195,10 +197,6 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                 tmpl.data = data
             self._initialized = True
 
-        # Set template accelerator use
-        for tmpl in self.templates:
-            tmpl.use_accel = use_accel
-
         # Set the data we are using for this execution
         for tmpl in self.templates:
             tmpl.det_data = self.det_data
@@ -221,22 +219,28 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                     raise RuntimeError(msg)
 
             if self.amplitudes not in data:
-                # The output template amplitudes do not yet exist.  Create these with
-                # all zero values.
+                # The output template amplitudes do not yet exist.
+                # Create these with all zero values.
                 data[self.amplitudes] = AmplitudesMap()
                 for tmpl in self.templates:
                     data[self.amplitudes][tmpl.name] = tmpl.zeros()
                 if use_accel:
+                    # We are running on the accelerator, so our output data must exist
+                    # on the device and will be used there.
                     data[self.amplitudes].accel_create()
+                    data[self.amplitudes].accel_update_device()
+            elif use_accel and (not data[self.amplitudes].accel_exists()):
+                # Our templates are using the accelerator to accumulate amplitudes.
+                # Ensure that the amplitudes are on the device.
+                data[self.amplitudes].accel_create()
+                data[self.amplitudes].accel_update_device()
             for d in all_dets:
                 for tmpl in self.templates:
                     log.verbose(f"TemplateMatrix {d} project_signal {tmpl.name}")
-                    tmpl.project_signal(d, data[self.amplitudes][tmpl.name])
+                    tmpl.project_signal(d, data[self.amplitudes][tmpl.name], use_accel=use_accel, **kwargs)
         else:
             if self.amplitudes not in data:
-                msg = "Template amplitudes '{}' do not exist in data".format(
-                    self.amplitudes
-                )
+                msg = f"Template amplitudes '{self.amplitudes}' do not exist in data"
                 log.error(msg)
                 raise RuntimeError(msg)
             # Ensure that our output detector data exists in each observation
@@ -257,22 +261,24 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                 for d in dets:
                     ob.detdata[self.det_data][d, :] = 0
                 if use_accel:
-                    if not exists and not ob.detdata.accel_present(self.det_data):
+                    if not ob.detdata.accel_exists(self.det_data):
                         ob.detdata.accel_create(self.det_data)
 
             for d in all_dets:
                 for tmpl in self.templates:
                     log.verbose(f"TemplateMatrix {d} add to signal {tmpl.name}")
-                    tmpl.add_to_signal(d, data[self.amplitudes][tmpl.name])
+                    tmpl.add_to_signal(d, data[self.amplitudes][tmpl.name], use_accel=use_accel, **kwargs)
         return
 
     def _finalize(self, data, use_accel=False, **kwargs):
         if self.transpose:
-            # Synchronize the result
+            # move amplitudes to host as sync is CPU only
             if use_accel:
                 data[self.amplitudes].accel_update_host()
+            # Synchronize the result
             for tmpl in self.templates:
                 data[self.amplitudes][tmpl.name].sync()
+            # move amplitudes back to GPU as it is NOT finalize's job to move data to host
             if use_accel:
                 data[self.amplitudes].accel_update_device()
         # Set the internal initialization to False, so that we are ready to process
@@ -305,6 +311,32 @@ def _provides(self):
             prov["detdata"] = [self.det_data]
         return prov
 
+    def _implementations(self):
+        """
+        Find implementations supported by all the templates
+        """
+        implementations = {
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        }
+        for tmpl in self.templates:
+            implementations.intersection_update(tmpl.implementations())
+        return list(implementations)
+
+    def _supports_accel(self):
+        """
+        Returns True if all the templates are GPU compatible.
+        """
+        for tmpl in self.templates:
+            if not tmpl.supports_accel():
+                log = Logger.get()
+                msg = f"{self} does not support accel because of '{tmpl}'"
+                log.debug(msg)
+                return False
+        return True
+
 
 @trait_docs
 class SolveAmplitudes(Operator):
diff --git a/src/toast/ops/mapmaker_utils/CMakeLists.txt b/src/toast/ops/mapmaker_utils/CMakeLists.txt
new file mode 100644
index 000000000..0c44b10b0
--- /dev/null
+++ b/src/toast/ops/mapmaker_utils/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    mapmaker_utils.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/mapmaker_utils
+)
diff --git a/src/toast/ops/mapmaker_utils/__init__.py b/src/toast/ops/mapmaker_utils/__init__.py
new file mode 100644
index 000000000..1f09fec75
--- /dev/null
+++ b/src/toast/ops/mapmaker_utils/__init__.py
@@ -0,0 +1,12 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .mapmaker_utils import (
+    BuildHitMap,
+    BuildInverseCovariance,
+    BuildNoiseWeighted,
+    CovarianceAndHits,
+)
diff --git a/src/toast/ops/mapmaker_utils/kernels.py b/src/toast/ops/mapmaker_utils/kernels.py
new file mode 100644
index 000000000..8dfabc4af
--- /dev/null
+++ b/src/toast/ops/mapmaker_utils/kernels.py
@@ -0,0 +1,189 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import healpy as hp
+import numpy as np
+
+from ..._libtoast import build_noise_weighted as libtoast_build_noise_weighted
+from ..._libtoast import cov_accum_diag_hits as libtoast_cov_accum_diag_hits
+from ..._libtoast import cov_accum_diag_invnpp as libtoast_cov_accum_diag_invnpp
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import (
+    build_noise_weighted_numpy,
+    cov_accum_diag_hits_numpy,
+    cov_accum_diag_invnpp_numpy,
+)
+
+if use_accel_jax:
+    from .kernels_jax import (
+        build_noise_weighted_jax,
+        cov_accum_diag_hits_jax,
+        cov_accum_diag_invnpp_jax,
+    )
+
+
+@kernel(impl=ImplementationType.COMPILED, name="build_noise_weighted")
+def build_noise_weighted_compiled(*args, use_accel=False):
+    return libtoast_build_noise_weighted(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="cov_accum_diag_hits")
+def cov_accum_diag_hits_compiled(*args, use_accel=False):
+    return libtoast_cov_accum_diag_hits(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="cov_accum_diag_invnpp")
+def cov_accum_diag_invnpp_compiled(*args, use_accel=False):
+    return libtoast_cov_accum_diag_invnpp(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def build_noise_weighted(
+    global2local,
+    zmap,
+    pixel_index,
+    pixels,
+    weight_index,
+    weights,
+    data_index,
+    det_data,
+    flag_index,
+    det_flags,
+    det_scale,
+    det_flag_mask,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel=False,
+):
+    """Kernel for accumulating the noise weighted map.
+
+    Args:
+        global2local (array):  The mapping from global submap to local submap index.
+        zmap (array):  The local piece of the noise weighted map, indexed by submap,
+            then pixel, then value.
+        pixel_index (array):  The index into the detector pixel array for each
+            detector.
+        pixels (array):  The array of detector pixels for each sample.
+        weight_index (array):  The index into the weights array for each detector.
+        weights (array):  The array of I, Q, and U weights at each sample for each
+            detector.
+        data_index (array):  The index into the data array for each detector.
+        det_data (array):  The detector data at each sample for each detector.
+        flag_index (array):  The index into the flag array for each detector.
+        det_flags (array):  The detector flag at each sample for each detector.
+        det_scale (float):  Scale factor for each detector, applied before accumulating.
+        det_flag_mask (int):  The flag mask to apply to the detector flags.
+        intervals (array):  The array of sample intervals.
+        shared_flags (array):  The array of common flags for each sample.
+        shared_flag_mask (int):  The flag mask to apply.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return build_noise_weighted(
+        global2local,
+        zmap,
+        pixel_index,
+        pixels,
+        weight_index,
+        weights,
+        data_index,
+        det_data,
+        flag_index,
+        det_flags,
+        det_scale,
+        det_flag_mask,
+        intervals,
+        shared_flags,
+        shared_flag_mask,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def cov_accum_diag_invnpp(
+    nsub,
+    nsubpix,
+    nnz,
+    submap,
+    subpix,
+    weights,
+    scale,
+    invnpp,
+    use_accel=False,
+):
+    """Kernel for accumulating the inverse diagonal pixel noise covariance.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index
+            within the local map (i.e. including only locally stored submaps)
+        subpix (array, int64):  For each time domain sample, the pixel index
+            within the submap.
+        weights (array, float64):  The pointing matrix weights for each time
+            sample and map.
+        scale (float):  Optional scaling factor.
+        invnpp (array, float64):  The local buffer of diagonal inverse pixel
+            covariances, stored as the lower triangle for each pixel.
+
+    Returns:
+        None
+
+    """
+    return cov_accum_diag_invnpp(
+        nsub,
+        nsubpix,
+        nnz,
+        submap,
+        subpix,
+        weights,
+        scale,
+        invnpp,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def cov_accum_diag_hits(
+    nsub,
+    nsubpix,
+    nnz,
+    submap,
+    subpix,
+    hits,
+    use_accel=False,
+):
+    """Kernel for accumulating the hits.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index
+            within the local map (i.e. including only locally stored submaps)
+        subpix (array, int64):  For each time domain sample, the pixel index
+            within the submap.
+        hits (array, int64):  The local hitmap buffer to accumulate.
+
+    Returns:
+        None
+
+    """
+    return cov_accum_diag_hits(
+        nsub,
+        nsubpix,
+        nnz,
+        submap,
+        subpix,
+        hits,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
diff --git a/src/toast/ops/mapmaker_utils/kernels_jax.py b/src/toast/ops/mapmaker_utils/kernels_jax.py
new file mode 100644
index 000000000..c6998056e
--- /dev/null
+++ b/src/toast/ops/mapmaker_utils/kernels_jax.py
@@ -0,0 +1,367 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+from jax.experimental.maps import xmap as jax_xmap
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import ALL, INTERVALS_JAX, JaxIntervals
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import AlignedF64, AlignedI64, Logger
+
+
+def build_noise_weighted_interval(
+    global2local,
+    zmap,
+    pixel_index,
+    pixels,
+    weight_index,
+    weights,
+    data_index,
+    det_data,
+    flag_index,
+    det_flags,
+    det_scale,
+    det_flag_mask,
+    shared_flags,
+    shared_flag_mask,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+):
+    """
+    Process all the intervals as a block.
+
+    Args:
+        global2local (array, int): size n_global_submap
+        zmap (array, double): size n_local_submap*n_pix_submap*nnz
+        pixel_index (array, int): size n_det
+        pixels (array, int): size n_det*n_samp
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, double): The flat packed detectors weights for the specified mode (size n_det*n_samp*nnz)
+        data_index (array, int): size n_det
+        det_data (array, double): size n_det*n_samp
+        flag_index (array, int): size n_det
+        det_flags (array, uint8): size n_det*n_samp
+        det_scale (array, double): size n_det
+        det_flag_mask (uint8)
+        shared_flags (array, uint8): size n_samp
+        shared_flag_mask (uint8)
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+
+    Returns:
+        zmap
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"build_noise_weighted: jit-compiling.")
+
+    # should we use flags?
+    n_samp = pixels.shape[1]
+    use_det_flags = det_flags.shape[1] == n_samp
+    use_shared_flags = shared_flags.size == n_samp
+
+    # extract interval slices
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    pixels_interval = JaxIntervals.get(
+        pixels, (pixel_index, intervals)
+    )  # pixels[pixel_index,intervals]
+    weights_interval = JaxIntervals.get(
+        weights, (weight_index, intervals, ALL)
+    )  # weights[weight_index,intervals,:]
+    data_interval = JaxIntervals.get(
+        det_data, (data_index, intervals)
+    )  # det_data[data_index,intervals]
+
+    # setup det check
+    if use_det_flags:
+        det_flags_interval = JaxIntervals.get(
+            det_flags, (flag_index, intervals)
+        )  # det_flags[flag_index,intervals]
+        det_check = (det_flags_interval & det_flag_mask) == 0
+    else:
+        det_check = True
+    # setup shared check
+    if use_shared_flags:
+        shared_flags_interval = JaxIntervals.get(
+            shared_flags, intervals
+        )  # shared_flags[intervals]
+        shared_check = (shared_flags_interval & shared_flag_mask) == 0
+    else:
+        shared_check = True
+    # mask to identify valid samples
+    valid_samples = (pixels_interval >= 0) & det_check & shared_check
+
+    # computes the update to add to zmap
+    scaled_data = data_interval * det_scale[:, jnp.newaxis, jnp.newaxis]
+    update = jnp.where(
+        valid_samples[:, :, :, jnp.newaxis],  # if
+        weights_interval * scaled_data[:, :, :, jnp.newaxis],  # then
+        0.0,
+    )  # else
+
+    # computes the index in zmap
+    n_pix_submap = zmap.shape[1]
+    global_submap = pixels_interval // n_pix_submap
+    local_submap = global2local[global_submap]
+    isubpix = pixels_interval - global_submap * n_pix_submap
+
+    # masks padded value before applying the update
+    update_masked = jnp.where(
+        intervals.mask[jnp.newaxis, :, :, jnp.newaxis],
+        zmap[local_submap, isubpix, :],
+        update,
+    )
+
+    # updates zmap and returns
+    zmap = zmap.at[local_submap, isubpix, :].add(update_masked)
+    return zmap
+
+
+# jit compiling
+build_noise_weighted_interval = jax.jit(
+    build_noise_weighted_interval,
+    static_argnames=["det_flag_mask", "shared_flag_mask", "intervals_max_length"],
+    donate_argnums=[1],
+)  # donate zmap
+
+
+@kernel(impl=ImplementationType.JAX, name="build_noise_weighted")
+def build_noise_weighted_jax(
+    global2local,
+    zmap,
+    pixel_index,
+    pixels,
+    weight_index,
+    weights,
+    data_index,
+    det_data,
+    flag_index,
+    det_flags,
+    det_scale,
+    det_flag_mask,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel,
+):
+    """
+    Args:
+        global2local (array, int): size n_global_submap
+        zmap (array, double): size n_local_submap*n_pix_submap*nnz
+        pixel_index (array, int): size n_det
+        pixels (array, int): size n_det*n_samp
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, double): The flat packed detectors weights for the specified mode (size n_det*n_samp*nnz)
+        data_index (array, int): size n_det
+        det_data (array, double): size n_det*n_samp
+        flag_index (array, int): size n_det
+        det_flags (array, uint8): size n_det*n_samp
+        det_scale (array, double): size n_det
+        det_flag_mask (uint8)
+        intervals (array, Interval): The intervals to modify (size n_view)
+        shared_flags (array, uint8): size n_samp
+        shared_flag_mask (uint8)
+        use_accel (Bool): should we use the accelerator?
+
+    Returns:
+        None (the result is put in zmap).
+    """
+    # prepares inputs
+    if intervals.size == 0:
+        return  # deals with a corner case in tests
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    zmap_input = MutableJaxArray.to_array(zmap)
+    global2local = MutableJaxArray.to_array(global2local)
+    pixels = MutableJaxArray.to_array(pixels)
+    pixel_index = MutableJaxArray.to_array(pixel_index)
+    weights = MutableJaxArray.to_array(weights)
+    weight_index = MutableJaxArray.to_array(weight_index)
+    det_data = MutableJaxArray.to_array(det_data)
+    data_index = MutableJaxArray.to_array(data_index)
+    det_flags = MutableJaxArray.to_array(det_flags)
+    flag_index = MutableJaxArray.to_array(flag_index)
+    det_scale = MutableJaxArray.to_array(det_scale)
+    shared_flags = MutableJaxArray.to_array(shared_flags)
+
+    # runs computation
+    zmap[:] = build_noise_weighted_interval(
+        global2local,
+        zmap_input,
+        pixel_index,
+        pixels,
+        weight_index,
+        weights,
+        data_index,
+        det_data,
+        flag_index,
+        det_flags,
+        det_scale,
+        det_flag_mask,
+        shared_flags,
+        shared_flag_mask,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+    )
+
+
+def cov_accum_diag_hits_inner(nsubpix, submap, subpix, hits):
+    """
+    Args:
+        nsubpix (int):  The number of pixels in each submap.
+        submap (array, int64):  For each time domain sample, the submap index within the local map (i.e. including only locally stored submaps) (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index within the submap (size nsamp).
+        hits (array, int64):  The local hitmap buffer to accumulate (size ???).
+
+    Returns:
+        hits
+    """
+    # computes update
+    added_value = jnp.where((submap >= 0) & (subpix >= 0), 1, 0)
+
+    # updates hits
+    hits = jnp.reshape(hits, newshape=(-1, nsubpix))
+    hits = hits.at[submap, subpix].add(added_value)
+    return hits.ravel()
+
+
+# jit compiling
+cov_accum_diag_hits_inner = jax.jit(
+    cov_accum_diag_hits_inner, static_argnames=["nsubpix"], donate_argnums=[3]
+)
+
+
+@kernel(impl=ImplementationType.JAX, name="cov_accum_diag_hits")
+def cov_accum_diag_hits_jax(nsub, nsubpix, nnz, submap, subpix, hits):
+    """
+    Accumulate hit map.
+    This uses a pointing matrix to accumulate the local pieces of the hit map.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index within the local map (i.e. including only locally stored submaps) (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index within the submap (size nsamp).
+        hits (array, int64):  The local hitmap buffer to accumulate (size ???).
+
+    Returns:
+        None (result is put in hits).
+    """
+    submap_input = MutableJaxArray.to_array(submap)
+    subpix_input = MutableJaxArray.to_array(subpix)
+    hits_input = MutableJaxArray.to_array(hits)
+    # AlignedI64 cannot be modified using numpy-style slices
+    # useful when running on CPU
+    if isinstance(hits, AlignedI64):
+        hits = hits_input
+    # run kernel
+    hits[:] = cov_accum_diag_hits_inner(nsubpix, submap_input, subpix_input, hits_input)
+
+
+def cov_accum_diag_invnpp_inner(nsubpix, nnz, submap, subpix, weights, scale, invnpp):
+    """
+    Args:
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index
+            within the local map (i.e. including only locally stored submaps)
+            (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index
+            within the submap (size nsamp).
+        weights (array, float64):  The pointing matrix weights for each time
+            sample and map (shape nw*nnz).
+        scale (float):  Optional scaling factor.
+        invnpp (array, float64):  The local buffer of diagonal inverse pixel
+            covariances, stored as the lower triangle for each pixel (shape ?*nsubpix*block with block=(nnz * (nnz + 1))/2).
+
+    Returns:
+        invnpp.
+    """
+    # reshape data
+    block = (nnz * (nnz + 1)) // 2
+    weights = jnp.reshape(weights, newshape=(-1, nnz))
+    invnpp = jnp.reshape(invnpp, newshape=(-1, nsubpix, block))
+
+    # converts flat index (i_block) back to index into upper triangular matrix of side nnz
+    # you can rederive the equations by knowing that i_block = col + row*nnz + row(row+1)/2
+    # then assuming col=row (close enough since row <= col < nnz) and rounding down to get row
+    i_block = jnp.arange(start=0, stop=block)
+    row = (2 * nnz + 1 - jnp.sqrt((2 * nnz + 1) ** 2 - 8 * i_block)).astype(int) // 2
+    col = i_block + (row * (row + 1)) // 2 - row * nnz
+
+    # computes mask
+    # newaxis are there to make dimenssion compatible with added_value
+    submap_2D = submap[..., jnp.newaxis]
+    subpix_2D = subpix[..., jnp.newaxis]
+    valid_index = (submap_2D >= 0) & (subpix_2D >= 0)
+
+    # updates invnpp
+    added_value = weights[:, col] * weights[:, row] * scale
+    masked_added_value = jnp.where(valid_index, added_value, 0.0)
+    invnpp = invnpp.at[submap, subpix, :].add(masked_added_value)
+
+    return invnpp.ravel()
+
+
+# jit compiling
+cov_accum_diag_invnpp_inner = jax.jit(
+    cov_accum_diag_invnpp_inner,
+    static_argnames=["nsubpix", "nnz"],
+    donate_argnums=[6],
+)
+
+
+@kernel(impl=ImplementationType.JAX, name="cov_accum_diag_invnpp")
+def cov_accum_diag_invnpp_jax(
+    nsub, nsubpix, nnz, submap, subpix, weights, scale, invnpp
+):
+    """
+    Accumulate block diagonal noise covariance.
+    This uses a pointing matrix to accumulate the local pieces
+    of the inverse diagonal pixel covariance.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index
+            within the local map (i.e. including only locally stored submaps)
+            (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index
+            within the submap (size nsamp).
+        weights (array, float64):  The pointing matrix weights for each time
+            sample and map (shape nw*nnz).
+        scale (float):  Optional scaling factor.
+        invnpp (array, float64):  The local buffer of diagonal inverse pixel
+            covariances, stored as the lower triangle for each pixel (shape ?*nsubpix*block with block=(nnz * (nnz + 1))/2).
+
+    Returns:
+        None (stores the result in invnpp).
+    """
+    submap_input = MutableJaxArray.to_array(submap)
+    subpix_input = MutableJaxArray.to_array(subpix)
+    weights_input = MutableJaxArray.to_array(weights)
+    invnpp_input = MutableJaxArray.to_array(invnpp)
+    # AlignedF64 cannot be modified using numpy-style slices
+    # useful when running on CPU
+    if isinstance(invnpp, AlignedF64):
+        invnpp = invnpp_input
+    # run kernel
+    invnpp[:] = cov_accum_diag_invnpp_inner(
+        nsubpix, nnz, submap_input, subpix_input, weights_input, scale, invnpp_input
+    )
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_mapmaker_utils"); toast.tests.run("covariance");'
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_sim_tod_conviqt", "ops_mapmaker_utils", "ops_mapmaker_binning", "ops_sim_tod_dipole", "ops_demodulate");'
diff --git a/src/toast/ops/mapmaker_utils/kernels_numpy.py b/src/toast/ops/mapmaker_utils/kernels_numpy.py
new file mode 100644
index 000000000..52da48c4d
--- /dev/null
+++ b/src/toast/ops/mapmaker_utils/kernels_numpy.py
@@ -0,0 +1,167 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="build_noise_weighted")
+def build_noise_weighted_numpy(
+    global2local,
+    zmap,
+    pixel_index,
+    pixels,
+    weight_index,
+    weights,
+    data_index,
+    det_data,
+    flag_index,
+    det_flags,
+    det_scale,
+    det_flag_mask,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel,
+):
+    """
+    Args:
+        global2local (array, int): size n_global_submap
+        zmap (array, double): size n_local_submap*n_pix_submap*nnz
+        pixel_index (array, int): size n_det
+        pixels (array, int): size ???*n_samp
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, double): The flat packed detectors weights for the specified mode (size ???*n_samp*nnz)
+        data_index (array, int): size n_det
+        det_data (array, double): size ???*n_samp
+        flag_index (array, int): size n_det
+        det_flags (array, uint8): size ???*n_samp or 1*1
+        det_scale (array, double): size n_det
+        det_flag_mask (uint8)
+        intervals (array, Interval): The intervals to modify (size n_view)
+        shared_flags (array, uint8): size n_samp
+        shared_flag_mask (uint8)
+        use_accel (Bool): should we use the accelerator?
+    Returns:
+        None (the result is put in zmap).
+    """
+    # should we use flags?
+    n_samp = pixels.shape[1]
+    skip_det_flags = det_flags.shape[1] != n_samp
+    skip_shared_flags = shared_flags.size != n_samp
+
+    # iterates on detectors and intervals
+    n_det = data_index.size
+    n_pix_submap = zmap.shape[1]
+    for idet in range(n_det):
+        p_index = pixel_index[idet]
+        w_index = weight_index[idet]
+        f_index = flag_index[idet]
+        d_index = data_index[idet]
+        det_scale_det = det_scale[idet]
+        for interval in intervals:
+            # get interval slices
+            interval_start = interval.first
+            interval_end = interval.last + 1
+            data_samples = det_data[d_index, interval_start:interval_end]
+            pixel_samples = pixels[p_index, interval_start:interval_end]
+            weights_sample = weights[w_index, interval_start:interval_end, :]
+            det_flags_samples = det_flags[f_index, interval_start:interval_end]
+            shared_flags_samples = shared_flags[interval_start:interval_end]
+            # keeps only good samples
+            det_check = skip_det_flags or ((det_flags_samples & det_flag_mask) == 0)
+            shared_check = skip_shared_flags or (
+                (shared_flags_samples & shared_flag_mask) == 0
+            )
+            good_samples = (pixel_samples >= 0) & (det_check & shared_check)
+            data_samples = data_samples[good_samples]
+            pixel_samples = pixel_samples[good_samples]
+            weights_sample = weights_sample[good_samples]
+            # computes the indices in zmap
+            global_submap = pixel_samples // n_pix_submap
+            local_submap = global2local[global_submap]
+            isubpix = pixel_samples - global_submap * n_pix_submap
+            # accumulates
+            # cannot use += due to duplicated indices, np.add being atomic
+            scaled_data = data_samples * det_scale_det
+            np.add.at(
+                zmap,
+                (local_submap, isubpix),
+                scaled_data[:, np.newaxis] * weights_sample,
+            )
+
+
+@kernel(impl=ImplementationType.NUMPY, name="cov_acccum_diag_hits")
+def cov_accum_diag_hits_numpy(nsub, nsubpix, nnz, submap, subpix, hits):
+    """
+    Accumulate hit map.
+    This uses a pointing matrix to accumulate the local pieces of the hit map.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index within the local map (i.e. including only locally stored submaps) (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index within the submap (size nsamp).
+        hits (array, int64):  The local hitmap buffer to accumulate (size ???).
+
+    Returns:
+        None (result is put in hits).
+    """
+    nsamp = submap.size
+    hits = np.asarray(hits).reshape((-1, nsubpix))
+
+    for i_samp in range(nsamp):
+        isubmap = submap[i_samp]
+        ipix = subpix[i_samp]
+        if (isubmap >= 0) and (ipix >= 0):
+            hits[isubmap, ipix] += 1
+
+
+@kernel(impl=ImplementationType.NUMPY, name="cov_accum_diag_invnpp")
+def cov_accum_diag_invnpp_numpy(
+    nsub, nsubpix, nnz, submap, subpix, weights, scale, invnpp
+):
+    """
+    Accumulate block diagonal noise covariance.
+    This uses a pointing matrix to accumulate the local pieces
+    of the inverse diagonal pixel covariance.
+
+    Args:
+        nsub (int):  The number of locally stored submaps.
+        nsubpix (int):  The number of pixels in each submap.
+        nnz (int):  The number of non-zeros in each row of the pointing matrix.
+        submap (array, int64):  For each time domain sample, the submap index
+            within the local map (i.e. including only locally stored submaps)
+            (size nsamp)
+        subpix (array, int64):  For each time domain sample, the pixel index
+            within the submap (size nsamp).
+        weights (array, float64):  The pointing matrix weights for each time
+            sample and map (shape nw*nnz).
+        scale (float):  Optional scaling factor.
+        invnpp (array, float64):  The local buffer of diagonal inverse pixel
+            covariances, stored as the lower triangle for each pixel (shape ?*nsubpix*block with block=(nnz * (nnz + 1))/2).
+
+    Returns:
+        None (stores the result in invnpp).
+    """
+    nsamp = submap.size
+    block = (nnz * (nnz + 1)) // 2
+    weights = np.asarray(weights).reshape((-1, nnz))
+    invnpp = np.asarray(invnpp).reshape((-1, nsubpix, block))
+
+    for i_samp in range(nsamp):
+        isubmap = submap[i_samp]
+        ipix = subpix[i_samp]
+        if (isubmap >= 0) and (ipix >= 0):
+            for i_block in range(block):
+                # converts i_block back to index into upper triangular matrix of side nnz
+                # you can rederive the equations by knowing that i_block = col + row*nnz + row(row+1)/2
+                # then assuming col=row (close enough since row <= col < nnz) and rounding down to get row
+                row = int(2 * nnz + 1 - np.sqrt((2 * nnz + 1) ** 2 - 8 * i_block)) // 2
+                col = i_block + (row * (row + 1)) // 2 - row * nnz
+                invnpp[isubmap, ipix, i_block] += (
+                    weights[i_samp, col] * weights[i_samp, row] * scale
+                )
diff --git a/src/toast/ops/mapmaker_utils.py b/src/toast/ops/mapmaker_utils/mapmaker_utils.py
similarity index 89%
rename from src/toast/ops/mapmaker_utils.py
rename to src/toast/ops/mapmaker_utils/mapmaker_utils.py
index 3bf92934e..a35711cda 100644
--- a/src/toast/ops/mapmaker_utils.py
+++ b/src/toast/ops/mapmaker_utils/mapmaker_utils.py
@@ -6,23 +6,18 @@
 import traitlets
 from astropy import units as u
 
-from .._libtoast import (
-    build_noise_weighted,
-    cov_accum_diag_hits,
-    cov_accum_diag_invnpp,
-    cov_accum_zmap,
-)
-from ..covariance import covariance_invert
-from ..mpi import MPI
-from ..observation import default_values as defaults
-from ..pixels import PixelData, PixelDistribution
-from ..timing import function_timer
-from ..traits import Bool, Float, Instance, Int, Unicode, Unit, trait_docs
-from ..utils import Logger, unit_conversion
-from .delete import Delete
-from .operator import Operator
-from .pipeline import Pipeline
-from .pointing import BuildPixelDistribution
+from ...accelerator import ImplementationType
+from ...covariance import covariance_invert
+from ...mpi import MPI
+from ...observation import default_values as defaults
+from ...pixels import PixelData
+from ...timing import function_timer
+from ...traits import Bool, Float, Instance, Int, Unicode, Unit, UseEnum, trait_docs
+from ...utils import Logger, unit_conversion
+from ..operator import Operator
+from ..pipeline import Pipeline
+from ..pointing import BuildPixelDistribution
+from .kernels import build_noise_weighted, cov_accum_diag_hits, cov_accum_diag_invnpp
 
 
 @trait_docs
@@ -102,9 +97,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.pixel_dist is None:
             raise RuntimeError(
                 "You must set the 'pixel_dist' trait before calling exec()"
@@ -185,6 +183,8 @@ def _exec(self, data, detectors=None, **kwargs):
                         local_sm.astype(np.int64),
                         local_pix.astype(np.int64),
                         hits.raw,
+                        impl=implementation,
+                        use_accel=use_accel,
                     )
         return
 
@@ -200,7 +200,7 @@ def _requires(self):
         req = {
             "global": [self.pixel_dist],
             "shared": list(),
-            "detdata": [self.pixels, self.weights],
+            "detdata": [self.pixels],
             "intervals": list(),
         }
         if self.shared_flags is not None:
@@ -215,6 +215,18 @@ def _provides(self):
         prov = {"global": [self.hits]}
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
+    def _supports_accel(self):
+        # NOTE: the kernels called do not follow the proper pattern yet
+        return False
+
 
 @trait_docs
 class BuildInverseCovariance(Operator):
@@ -305,9 +317,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.pixel_dist is None:
             raise RuntimeError(
                 "You must set the 'pixel_dist' trait before calling exec()"
@@ -453,6 +468,8 @@ def _exec(self, data, detectors=None, **kwargs):
                         wview[det].reshape(-1),
                         detweight.to_value(invcov_units),
                         invcov.raw,
+                        impl=implementation,
+                        use_accel=use_accel,
                     )
         return
 
@@ -484,6 +501,18 @@ def _provides(self):
         prov = {"global": [self.inverse_covariance]}
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
+    def _supports_accel(self):
+        # NOTE: the kernels called do not follow the proper pattern yet
+        return False
+
 
 @trait_docs
 class BuildNoiseWeighted(Operator):
@@ -562,8 +591,6 @@ class BuildNoiseWeighted(Operator):
         "alltoallv", help="Communication algorithm: 'allreduce' or 'alltoallv'"
     )
 
-    use_python = Bool(False, help="If True, use python implementation")
-
     @traitlets.validate("det_flag_mask")
     def _check_flag_mask(self, proposal):
         check = proposal["value"]
@@ -585,6 +612,9 @@ def __init__(self, **kwargs):
     def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.pixel_dist is None:
             raise RuntimeError(
                 "You must set the 'pixel_dist' trait before calling exec()"
@@ -600,9 +630,6 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         if self.det_data is None:
             raise RuntimeError("You must set the det_data trait before calling exec()")
 
-        if self.use_python and use_accel:
-            raise RuntimeError("Cannot use accelerator with pure python implementation")
-
         dist = data[self.pixel_dist]
         if data.comm.world_rank == 0:
             log.verbose(
@@ -719,56 +746,39 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
             else:
                 shared_flag_data = np.zeros(1, dtype=np.uint8)
 
-            if self.use_python:
-                self._py_build_noise_weighted(
-                    zmap,
-                    pix_indx,
-                    ob.detdata[self.pixels].data,
-                    weight_indx,
-                    ob.detdata[self.weights].data.reshape(
-                        (n_weight_dets, ob.n_local_samples, -1)
-                    ),
-                    data_indx,
-                    ob.detdata[self.det_data].data,
-                    flag_indx,
-                    flag_data,
-                    self.det_flag_mask,
-                    ob.intervals[self.view].data,
-                    shared_flag_data,
-                    self.shared_flag_mask,
-                    detweights,
-                )
-            else:
-                build_noise_weighted(
-                    zmap.distribution.global_submap_to_local.array(),
-                    zmap.data.reshape(
-                        (zmap.distribution.n_local_submap, -1, weight_nnz)
-                    ),
-                    pix_indx,
-                    ob.detdata[self.pixels].data,
-                    weight_indx,
-                    ob.detdata[self.weights].data.reshape(
-                        (n_weight_dets, ob.n_local_samples, -1)
-                    ),
-                    data_indx,
-                    ob.detdata[self.det_data].data,
-                    flag_indx,
-                    flag_data,
-                    detweights,
-                    self.det_flag_mask,
-                    ob.intervals[self.view].data,
-                    shared_flag_data,
-                    self.shared_flag_mask,
-                    use_accel,
-                )
-
+            # uses a separate variable in case the reshaping copies the data
+            zmap_data = zmap.data.reshape(
+                (zmap.distribution.n_local_submap, -1, weight_nnz)
+            )
+            build_noise_weighted(
+                zmap.distribution.global_submap_to_local,
+                zmap_data,
+                pix_indx,
+                ob.detdata[self.pixels].data,
+                weight_indx,
+                ob.detdata[self.weights].data.reshape(
+                    (n_weight_dets, ob.n_local_samples, -1)
+                ),
+                data_indx,
+                ob.detdata[self.det_data].data,
+                flag_indx,
+                flag_data,
+                detweights,
+                self.det_flag_mask,
+                ob.intervals[self.view].data,
+                shared_flag_data,
+                self.shared_flag_mask,
+                impl=implementation,
+                use_accel=use_accel,
+            )
+            zmap.data = zmap_data.reshape(zmap.data.shape)
         return
 
-    def _finalize(self, data, use_acc=False, **kwargs):
+    def _finalize(self, data, use_accel=False, **kwargs):
         if self.zmap in data:
             log = Logger.get()
             # We have called exec() at least once
-            if use_acc:
+            if use_accel:
                 log.verbose_rank(
                     f"Operator {self.name} finalize calling zmap update self",
                     comm=data.comm.comm_group,
@@ -778,7 +788,7 @@ def _finalize(self, data, use_acc=False, **kwargs):
                 data[self.zmap].sync_alltoallv()
             else:
                 data[self.zmap].sync_allreduce()
-            if use_acc:
+            if use_accel:
                 log.verbose_rank(
                     f"Operator {self.name} finalize calling zmap update device",
                     comm=data.comm.comm_group,
@@ -825,59 +835,17 @@ def _provides(self):
         }
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
     def _supports_accel(self):
         return True
 
-    def _py_build_noise_weighted(
-        self,
-        zmap,
-        pixel_indx,
-        pixel_data,
-        weight_indx,
-        weight_data,
-        det_indx,
-        det_data,
-        flag_indx,
-        flag_data,
-        flag_mask,
-        intr_data,
-        shared_flags,
-        shared_mask,
-        det_scale,
-    ):
-        """Internal python implementation for comparison tests."""
-        global2local = zmap.distribution.global_submap_to_local.array()
-        npix_submap = zmap.distribution.n_pix_submap
-        nnz = zmap.n_value
-        for idet in range(len(det_indx)):
-            didx = det_indx[idet]
-            pidx = pixel_indx[idet]
-            widx = weight_indx[idet]
-            fidx = flag_indx[idet]
-            for vw in intr_data:
-                samples = slice(vw.first, vw.last + 1, 1)
-                good = np.logical_and(
-                    ((flag_data[fidx][samples] & flag_mask) == 0),
-                    ((shared_flags[samples] & shared_mask) == 0),
-                )
-                pixel_buffer = pixel_data[pidx][samples]
-                det_buffer = det_data[didx][samples]
-                weight_buffer = weight_data[widx][samples]
-                global_submap = pixel_buffer[good] // npix_submap
-                submap_pix = pixel_buffer[good] - global_submap * npix_submap
-                local_submap = np.array(
-                    [global2local[x] for x in global_submap], dtype=np.int64
-                )
-                tempdata = np.multiply(
-                    weight_buffer[good],
-                    np.multiply(det_scale[idet], det_buffer[good])[:, np.newaxis],
-                )
-                np.add.at(
-                    zmap.data,
-                    (local_submap, submap_pix),
-                    tempdata,
-                )
-
 
 @trait_docs
 class CovarianceAndHits(Operator):
diff --git a/src/toast/ops/noise_estimation.py b/src/toast/ops/noise_estimation.py
index e60bea4c3..911eb5617 100644
--- a/src/toast/ops/noise_estimation.py
+++ b/src/toast/ops/noise_estimation.py
@@ -4,9 +4,6 @@
 
 import copy
 import os
-import re
-import warnings
-from time import time
 
 import astropy.io.fits as pf
 import numpy as np
@@ -14,23 +11,11 @@
 import traitlets
 from astropy import units as u
 
-from .. import qarray as qa
-from .._libtoast import filter_poly2D, filter_polynomial, subtract_mean, sum_detectors
-from ..intervals import interval_dtype
-from ..mpi import MPI, Comm, MPI_Comm, use_mpi
 from ..noise import Noise
 from ..observation import default_values as defaults
 from ..timing import function_timer
 from ..traits import Bool, Instance, Int, List, Quantity, Tuple, Unicode, trait_docs
-from ..utils import (
-    AlignedF64,
-    AlignedU8,
-    Environment,
-    GlobalTimers,
-    Logger,
-    Timer,
-    dtype_to_aligned,
-)
+from ..utils import Logger, Timer
 from .arithmetic import Combine
 from .copy import Copy
 from .delete import Delete
diff --git a/src/toast/ops/noise_weight/CMakeLists.txt b/src/toast/ops/noise_weight/CMakeLists.txt
new file mode 100644
index 000000000..6dcd12e08
--- /dev/null
+++ b/src/toast/ops/noise_weight/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    noise_weight.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/noise_weight
+)
diff --git a/src/toast/ops/noise_weight/__init__.py b/src/toast/ops/noise_weight/__init__.py
new file mode 100644
index 000000000..87f7635cc
--- /dev/null
+++ b/src/toast/ops/noise_weight/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .noise_weight import NoiseWeight
diff --git a/src/toast/ops/noise_weight/kernels.py b/src/toast/ops/noise_weight/kernels.py
new file mode 100644
index 000000000..8bbfca59f
--- /dev/null
+++ b/src/toast/ops/noise_weight/kernels.py
@@ -0,0 +1,50 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+# from ..._libtoast import noise_weight as libtoast_noise_weight
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import noise_weight_numpy
+
+if use_accel_jax:
+    from .kernels_jax import noise_weight_jax
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def noise_weight(
+    det_data,
+    det_data_index,
+    intervals,
+    detector_weights,
+    use_accel=False,
+):
+    """Kernel for applying noise weights to detector timestreams.
+
+    Args:
+        det_data (array):  The detector data at each sample for each detector.
+        det_data_index (array):  The index into the data array for each detector.
+        intervals (array):  The array of sample intervals.
+        detector_weights (array):  Array of noise weights for each detector.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return noise_weight(
+        det_data,
+        det_data_index,
+        intervals,
+        detector_weights,
+        impl=ImplementationType.NUMPY,
+        use_accel=use_accel,
+    )
+
+
+# FIXME:  Add this to _libtoast
+
+# @kernel(impl=ImplementationType.COMPILED, name="noise_weight")
+# def noise_weight_compiled(*args, use_accel=False):
+#     return libtoast_noise_weight(*args, use_accel)
diff --git a/src/toast/ops/noise_weight/kernels_jax.py b/src/toast/ops/noise_weight/kernels_jax.py
new file mode 100644
index 000000000..053c5c736
--- /dev/null
+++ b/src/toast/ops/noise_weight/kernels_jax.py
@@ -0,0 +1,100 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import INTERVALS_JAX, JaxIntervals
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def noise_weight_interval(
+    det_data,
+    det_data_index,
+    detector_weights,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+):
+    """
+    multiplies det_data by the weighs in detector_weights, applied to all intervals as a single block
+
+    Args:
+        det_data (array, float): size ???*n_samp
+        det_data_index (array, int): The indexes of the det_data (size n_det)
+        detector_weights (list, double): The weight to be used for each detcetor (size n_det)
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+
+    Returns:
+        det_data
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"noise_weight: jit-compiling.")
+
+    # turns detector_weights into a jax array
+    detector_weights = jnp.array(detector_weights)
+
+    # extract interval slice
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    det_data_interval = JaxIntervals.get(
+        det_data, (det_data_index, intervals)
+    )  # det_data[det_data_index, intervals]
+
+    # does the computation
+    new_det_data_interval = (
+        det_data_interval * detector_weights[:, jnp.newaxis, jnp.newaxis]
+    )
+
+    # updates results and returns
+    # det_data[det_data_index, intervals] = new_det_data_interval
+    det_data = JaxIntervals.set(
+        det_data, (det_data_index, intervals), new_det_data_interval
+    )
+    return det_data
+
+
+# jit compiling
+noise_weight_interval = jax.jit(
+    noise_weight_interval,
+    static_argnames=["intervals_max_length"],
+    donate_argnums=[0],
+)  # donates det_data
+
+
+@kernel(impl=ImplementationType.JAX, name="noise_weight")
+def noise_weight_jax(det_data, det_data_index, intervals, detector_weights, use_accel):
+    """
+    multiplies det_data by the weighs in detector_weights
+
+    Args:
+        det_data (array, float): size ???*n_samp
+        det_data_index (array, int): The indexes of the det_data (size n_det)
+        intervals (array, Interval): The intervals to modify (size n_view)
+        detector_weights (list, double): The weight to be used for each detcetor (size n_det)
+        use_accel (bool): should we use an accelerator
+
+    Returns:
+        None: det_data is updated in place
+    """
+    # prepares inputs
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    det_data_input = MutableJaxArray.to_array(det_data)
+    det_data_index = MutableJaxArray.to_array(det_data_index)
+
+    # performs computation and updates det_data in place
+    det_data[:] = noise_weight_interval(
+        det_data_input,
+        det_data_index,
+        detector_weights,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+    )
diff --git a/src/toast/ops/noise_weight/kernels_numpy.py b/src/toast/ops/noise_weight/kernels_numpy.py
new file mode 100644
index 000000000..88949a064
--- /dev/null
+++ b/src/toast/ops/noise_weight/kernels_numpy.py
@@ -0,0 +1,27 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="noise_weight")
+def noise_weight_numpy(
+    det_data,
+    det_data_index,
+    intervals,
+    detector_weights,
+    use_accel,
+):
+    # Iterates over detectors and intervals
+    n_det = det_data_index.size
+    for idet in range(n_det):
+        d_index = det_data_index[idet]
+        detector_weight = detector_weights[idet]
+        for interval in intervals:
+            interval_start = interval.first
+            interval_end = interval.last + 1
+            # Multiply by the weight
+            det_data[d_index, interval_start:interval_end] *= detector_weight
diff --git a/src/toast/ops/noise_weight.py b/src/toast/ops/noise_weight/noise_weight.py
similarity index 60%
rename from src/toast/ops/noise_weight.py
rename to src/toast/ops/noise_weight/noise_weight.py
index 231acf016..c37b340e8 100644
--- a/src/toast/ops/noise_weight.py
+++ b/src/toast/ops/noise_weight/noise_weight.py
@@ -5,11 +5,13 @@
 import numpy as np
 import traitlets
 
-from ..noise_sim import AnalyticNoise
-from ..timing import Timer, function_timer
-from ..traits import Int, Unicode, trait_docs
-from ..utils import Environment, Logger
-from .operator import Operator
+from ...accelerator import ImplementationType
+from ...noise_sim import AnalyticNoise
+from ...timing import Timer, function_timer
+from ...traits import Int, Unicode, UseEnum, trait_docs
+from ...utils import Environment, Logger
+from ..operator import Operator
+from .kernels import noise_weight
 
 
 @trait_docs
@@ -41,9 +43,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         for ob in data.obs:
             # Get the detectors we are using for this observation
             dets = ob.select_local_detectors(detectors)
@@ -62,17 +67,29 @@ def _exec(self, data, detectors=None, **kwargs):
                 )
                 raise RuntimeError(msg)
 
+            # Compute the noise for each detector (using the correct units)
             noise = ob[self.noise_model]
-
-            for vw in ob.view[self.view].detdata[self.det_data]:
-                for d in dets:
-                    # Get the detector weight from the noise model.
-                    detweight = noise.detector_weight(d).to(data_invcov_units)
-
-                    # Apply
-                    vw[d] *= detweight.value
-
+            detector_weights = [
+                noise.detector_weight(detector).to(data_invcov_units).value
+                for detector in dets
+            ]
+
+            # Multiply detectors by their respective noise weight
+            intervals = ob.intervals[self.view].data
+            det_data = ob.detdata[self.det_data].data
+            det_data_indx = ob.detdata[self.det_data].indices(dets)
+            noise_weight(
+                det_data,
+                det_data_indx,
+                intervals,
+                detector_weights,
+                impl=implementation,
+                use_accel=use_accel,
+            )
+
+            # Update the units of the output
             ob.detdata[self.det_data].update_units(data_output_units)
+
         return
 
     def _finalize(self, data, **kwargs):
@@ -90,3 +107,14 @@ def _requires(self):
 
     def _provides(self):
         return dict()
+
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
+    def _supports_accel(self):
+        return True
diff --git a/src/toast/ops/operator.py b/src/toast/ops/operator.py
index f1b4370b9..e862f9655 100644
--- a/src/toast/ops/operator.py
+++ b/src/toast/ops/operator.py
@@ -28,11 +28,6 @@ def exec(self, data, detectors=None, use_accel=False, **kwargs):
         If a list of detectors is specified, only process these detectors.  Any extra
         kwargs are passed to the derived class internal method.
 
-        Accelerator use:  If the derived class supports OpenACC and all the required
-        data objects exist on the device, then the `_exec()` method will be called
-        with the "use_accel=True" option.  Any operator that returns "True" from its
-        _supports_accel() method should also accept the "use_accel" keyword argument.
-
         Args:
             data (toast.Data):  The distributed data.
             detectors (list):  A list of detector names or indices.  If None, this
@@ -169,18 +164,6 @@ def provides(self):
                 prov[key] = list()
         return prov
 
-    def _supports_accel(self):
-        return False
-
-    def supports_accel(self):
-        """Query whether the operator supports OpenACC
-
-        Returns:
-            (bool):  True if the operator can use OpenACC, else False.
-
-        """
-        return self._supports_accel()
-
     def accel_have_requires(self, data):
         # Helper function to determine if all requirements are met to use accelerator
         # for a data object.
@@ -213,11 +196,9 @@ def accel_have_requires(self, data):
                     msg = f"{self.name}:  obs {ob.name}, intervals {key} is on device"
                 log.verbose(msg)
         if all_present:
-            log.verbose(f"{self.name}:  obs {ob.name} all required inputs on device")
+            log.verbose(f"{self.name}:  all required inputs on device")
         else:
-            log.verbose(
-                f"{self.name}:  obs {ob.name} some required inputs not on device"
-            )
+            log.verbose(f"{self.name}:  some required inputs not on device")
         return all_present
 
     @classmethod
@@ -276,3 +257,10 @@ def translate(cls, props):
         if "API" in props:
             del props["API"]
         return props
+
+    def __str__(self):
+        """
+        Converts the operator into a short human readable string
+        returns only its class name (for example: 'Operator')
+        """
+        return str(self.__class__).split(".")[-1][:-2]
diff --git a/src/toast/ops/pipeline.py b/src/toast/ops/pipeline.py
index 5fe5b60d0..5b07967af 100644
--- a/src/toast/ops/pipeline.py
+++ b/src/toast/ops/pipeline.py
@@ -6,7 +6,7 @@
 
 import traitlets
 
-from ..accelerator import use_accel_jax, use_accel_omp
+from ..accelerator import ImplementationType, accel_enabled
 from ..data import Data
 from ..timing import function_timer
 from ..traits import Int, List, trait_docs
@@ -40,9 +40,9 @@ class Pipeline(Operator):
     def _check_detsets(self, proposal):
         detsets = proposal["value"]
         if len(detsets) == 0:
-            raise traitlets.TraitError(
-                "detector_sets must be a list with at least one entry ('ALL' and 'SINGLE' are valid entries)"
-            )
+            msg = "detector_sets must be a list with at least one entry "
+            msg += "('ALL' and 'SINGLE' are valid entries)"
+            raise traitlets.TraitError(msg)
         for dset in detsets:
             if (dset != "ALL") and (dset != "SINGLE"):
                 # Not a built-in name, must be an actual list of detectors
@@ -82,73 +82,50 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                 "Pipeline has no operators, nothing to do", comm=data.comm.comm_world
             )
             return
-
-        # By default, if the calling code passed use_accel=True, then we assume the
-        # data staging is being handled at a higher level.
+        
+        # There are 2 scenarios we handle here:  If the use_accel trait has already
+        # been set to True, it means that the calling code has already queried the
+        # supports_accel() method and it returned True, and the calling code has
+        # ensured that all requirements are staged to the accelerator.  If the
+        # use_accel trait is False, then we still may be able to run on the accelerator
+        # if our operators support it and if we stage the data ourselves.
+
+        # This is used to track whether we are handling the data staging internally.
+        # If this is set to True in the finalize method, then we copy data back out
+        # and restore our state.
         self._staged_accel = False
 
         if not use_accel:
             # The calling code determined that we do not have all the data present to
-            # use the accelerator.  HOWEVER, if our operators support it, we can stage
-            # the data to and from the device.
+            # use the accelerator.  However, if all of our operators support it, we 
+            # can stage the data to and from the device.
             if detectors is None:
                 comp_dets = set(data.all_local_detectors(selection=None))
             else:
                 comp_dets = set(detectors)
-            if (use_accel_omp or use_accel_jax) and self.supports_accel():
+            if accel_enabled() and self.supports_accel():
                 # All our operators support it.
-                msg = "Pipeline operators {}".format(
-                    ", ".join([x.name for x in self.operators])
-                )
-                msg += " all support accelerators, data to be staged: "
-                msg += f"{self.requires()}"
+                msg = f"{self} fully supports accelerators, data to "
+                msg += f"be staged: {self.requires()}"
                 log.verbose_rank(msg, comm=data.comm.comm_world)
 
-                interm = self._get_intermediate()
-                log.verbose(f"intermediate = {interm}")
-                for ob in data.obs:
-                    for obj in interm["detdata"]:
-                        if obj in ob.detdata and not ob.detdata.accel_exists(obj):
-                            # The data exists on the host from a previous call, but is
-                            # not on the device.  Delete the host copy so that it will
-                            # be re-created consistently.
-                            msg = f"Pipeline intermediate detdata '{obj}' in "
-                            msg += f"observation '{ob.name}' exists on "
-                            msg += f"the host but not the device.  Deleting."
-                            log.verbose_rank(msg, comm=data.comm.comm_group)
-                            del ob.detdata[obj]
-                    for obj in interm["shared"]:
-                        if obj in ob.shared and not ob.shared.accel_exists(obj):
-                            msg = f"Pipeline intermediate shared data '{obj}' in "
-                            msg += f"observation '{ob.name}' exists on "
-                            msg += f"the host but not the device.  Deleting."
-                            log.verbose_rank(msg, comm=data.comm.comm_group)
-                            del ob.shared[obj]
-                    for obj in interm["intervals"]:
-                        if obj in ob.intervals and not ob.intervals.accel_exists(obj):
-                            msg = f"Pipeline intermediate intervals '{obj}' in "
-                            msg += f"observation '{ob.name}' exists on "
-                            msg += f"the host but not the device.  Deleting."
-                            log.verbose_rank(msg, comm=data.comm.comm_group)
-                            del ob.intervals[obj]
-
-                data.accel_create(self.requires())
-                data.accel_update_device(self.requires())
-
-                use_accel = True
+                # Deletes leftover intermediate values
+                self._delete_intermediates(data)
+                
+                # Send the requirements to the device
+                self._stage_requirements_to_device(data)
                 self._staged_accel = True
+                use_accel = True
 
         if len(data.obs) == 0:
             # No observations for this group
-            msg = f"Pipeline data, group {data.comm.group} has no observations."
+            msg = f"{self} data, group {data.comm.group} has no observations."
             log.verbose_rank(msg, comm=data.comm.comm_group)
 
         if len(self.detector_sets) == 1 and self.detector_sets[0] == "ALL":
             # Run the operators with all detectors at once
             for op in self.operators:
-                msg = "{} Pipeline calling operator '{}' exec() with ALL dets".format(
-                    pstr, op.name
-                )
+                msg = f"{pstr} {self} calling operator '{op.name}' exec() with ALL dets"
                 log.verbose(msg)
                 op.exec(data, detectors=None, use_accel=use_accel)
         elif len(self.detector_sets) == 1 and self.detector_sets[0] == "SINGLE":
@@ -156,12 +133,10 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
             all_local_dets = data.all_local_detectors(selection=detectors)
             # Run operators one detector at a time
             for det in all_local_dets:
-                msg = "{} Pipeline SINGLE detector {}".format(pstr, det)
+                msg = f"{pstr} {self} SINGLE detector {det}"
                 log.verbose(msg)
                 for op in self.operators:
-                    msg = "{} Pipeline   calling operator '{}' exec()".format(
-                        pstr, op.name
-                    )
+                    msg = f"{pstr} {self} calling operator '{op.name}' exec()"
                     log.verbose(msg)
                     op.exec(data, detectors=[det], use_accel=use_accel)
         else:
@@ -177,48 +152,96 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                 if len(selected_set) == 0:
                     # Nothing in this detector set is being used, skip it
                     continue
-                msg = "{} Pipeline detector set {}".format(pstr, selected_set)
+                msg = f"{pstr} {self} detector set {selected_set}"
                 log.verbose(msg)
                 for op in self.operators:
-                    msg = "{} Pipeline   calling operator '{}' exec()".format(
-                        pstr, op.name
-                    )
+                    msg = f"{pstr} {self} calling operator '{op.name}' exec()"
                     log.verbose(msg)
                     op.exec(data, detectors=selected_set, use_accel=use_accel)
-
         return
 
+    @function_timer
+    def _delete_intermediates(self, data):
+        """
+        Deals with data existing on the host from a previous call, but
+        not on the device.  Delete the host copy so that it will
+        be re-created consistently.
+        """
+        log = Logger.get()
+        interm = self._get_intermediate()
+        log.verbose(f"intermediate = {interm}")
+        for ob in data.obs:
+            for obj in interm["detdata"]:
+                if obj in ob.detdata and not ob.detdata.accel_exists(obj):
+                    msg = f"Pipeline intermediate detdata '{obj}' in "
+                    msg += f"observation '{ob.name}' exists on "
+                    msg += f"the host but not the device.  Deleting."
+                    log.verbose_rank(msg, comm=data.comm.comm_group)
+                    del ob.detdata[obj]
+            for obj in interm["shared"]:
+                if obj in ob.shared and not ob.shared.accel_exists(obj):
+                    msg = f"Pipeline intermediate shared data '{obj}' in "
+                    msg += f"observation '{ob.name}' exists on "
+                    msg += f"the host but not the device.  Deleting."
+                    log.verbose_rank(msg, comm=data.comm.comm_group)
+                    del ob.shared[obj]
+            for obj in interm["intervals"]:
+                if obj in ob.intervals and not ob.intervals.accel_exists(obj):
+                    msg = f"Pipeline intermediate intervals '{obj}' in "
+                    msg += f"observation '{ob.name}' exists on "
+                    msg += f"the host but not the device.  Deleting."
+                    log.verbose_rank(msg, comm=data.comm.comm_group)
+                    del ob.intervals[obj]
+
+    @function_timer
+    def _stage_requirements_to_device(self, data):
+        """Move required data to the device"""
+        requires = self.requires()
+        data.accel_create(requires)
+        data.accel_update_device(requires)
+
     @function_timer
     def _finalize(self, data, use_accel=False, **kwargs):
         log = Logger.get()
         result = list()
         pstr = f"Proc ({data.comm.world_rank}, {data.comm.group_rank})"
+        msg = f"{pstr} {self} finalize"
+        log.verbose(msg)
 
         # FIXME:  We need to clarify in documentation that if using the
         # accelerator in _finalize() to produce output products, these
         # outputs should remain on the device so that they can be copied
         # out at the end automatically.
 
-        if not use_accel and self._staged_accel:
-            use_accel = True
-
         if self.operators is not None:
             for op in self.operators:
-                msg = f"{pstr} Pipeline calling operator '{op.name}' finalize()"
                 log.verbose(msg)
-                result.append(op.finalize(data, use_accel=use_accel))
+                result.append(op.finalize(data, use_accel=use_accel, **kwargs))
 
         # Copy out from accelerator if we did the copy in.
         if self._staged_accel:
+            # Copy out the outputs to the CPU
             prov = self.provides()
-            msg = f"{pstr} Pipeline copying out accel data products: {prov}"
+            msg = f"{pstr} {self} copying out accel data outputs: {prov}"
             log.verbose(msg)
-            data.accel_update_host(self.provides())
+            data.accel_update_host(prov)
+            # Delete the intermediate products from the GPU.  Otherwise, they will
+            # get re-used by other pipelines despite still being on GPU.
+            interm = self._get_intermediate()
+            msg = f"{pstr} {self} deleting accel data intermediate outputs: {interm}"
+            log.verbose(msg)
+            data.accel_delete(interm)
+            # Delete the inputs.  Otherwise, they will get re-used by other pipelines
+            # despite still being on GPU.
+            req = self.requires()
+            msg = f"{pstr} {self} deleting accel data inputs: {req}"
+            log.verbose(msg)
+            data.accel_delete(req)
         return result
 
     def _requires(self):
         # Work through the operator list in reverse order and prune intermediate
-        # products.
+        # products (that will be provided by a previous operator).
         if self.operators is None:
             return dict()
         keys = ["global", "meta", "detdata", "shared", "intervals"]
@@ -237,7 +260,9 @@ def _requires(self):
         return req
 
     def _provides(self):
-        # Work through the operator list and prune intermediate products.
+        # Work through the operator list and prune intermediate products
+        # (that are be provided to an intermediate operator).
+        # FIXME could a final result also be used by an intermediate operator?
         if self.operators is None:
             return dict()
         keys = ["global", "meta", "detdata", "shared", "intervals"]
@@ -256,6 +281,7 @@ def _provides(self):
 
     def _get_intermediate(self):
         keys = ["global", "meta", "detdata", "shared", "intervals"]
+        # Full provide minus intermediate
         prov = self.provides()
         interm = {x: set() for x in keys}
         for op in self.operators:
@@ -263,14 +289,40 @@ def _get_intermediate(self):
             for k in keys:
                 if k in oprov:
                     interm[k] |= set(oprov[k])
+        # Deduce intermediate by subtraction
         for k in keys:
             interm[k] -= set(prov[k])
             interm[k] = list(interm[k])
         return interm
 
+    def _implementations(self):
+        """
+        Find implementations supported by all the operators
+        """
+        implementations = {
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        }
+        for op in self.operators:
+            implementations.intersection_update(op.implementations())
+        return list(implementations)
+
     def _supports_accel(self):
-        # This is a logical AND of our operators
+        """
+        Returns True if all the operators are GPU compatible.
+        """
         for op in self.operators:
             if not op.supports_accel():
+                log = Logger.get()
+                msg = f"{self} does not support accel because of '{op}'"
+                log.debug(msg)
                 return False
         return True
+
+    def __str__(self):
+        """
+        Converts the pipeline into a human-readable string.
+        """
+        return f"Pipeline{[str(op) for op in self.operators]}"
diff --git a/src/toast/ops/pixels_healpix/CMakeLists.txt b/src/toast/ops/pixels_healpix/CMakeLists.txt
new file mode 100644
index 000000000..953ab69c4
--- /dev/null
+++ b/src/toast/ops/pixels_healpix/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    pixels_healpix.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/pixels_healpix
+)
diff --git a/src/toast/ops/pixels_healpix/__init__.py b/src/toast/ops/pixels_healpix/__init__.py
new file mode 100644
index 000000000..615cb13a8
--- /dev/null
+++ b/src/toast/ops/pixels_healpix/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .pixels_healpix import PixelsHealpix
diff --git a/src/toast/ops/pixels_healpix/kernels.py b/src/toast/ops/pixels_healpix/kernels.py
new file mode 100644
index 000000000..781917c0f
--- /dev/null
+++ b/src/toast/ops/pixels_healpix/kernels.py
@@ -0,0 +1,74 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import healpy as hp
+import numpy as np
+
+from ... import qarray as qa
+from ..._libtoast import pixels_healpix as libtoast_pixels_healpix
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import pixels_healpix_numpy
+
+if use_accel_jax:
+    from .kernels_jax import pixels_healpix_jax
+
+
+@kernel(impl=ImplementationType.COMPILED, name="pixels_healpix")
+def pixels_healpix_compiled(*args, use_accel=False):
+    return libtoast_pixels_healpix(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def pixels_healpix(
+    quat_index,
+    quats,
+    shared_flags,
+    shared_flag_mask,
+    pixel_index,
+    pixels,
+    intervals,
+    hit_submaps,
+    n_pix_submap,
+    nside,
+    nest,
+    use_accel=False,
+):
+    """Kernel for computing healpix pixelization.
+
+    Args:
+        quat_index (array):  The index into the detector quaternion array for each
+            detector.
+        quats (array):  The array of detector quaternions for each sample.
+        shared_flags (array):  The array of common flags for each sample.
+        shared_flag_mask (int):  The flag mask to apply.
+        pixel_index (array):  The index into the detector pixel array for each
+            detector.
+        pixels (array):  The array of detector pixels for each sample.
+        intervals (array):  The array of sample intervals.
+        hit_submaps (array):  Array of bytes to set to 1 if the submap is hit
+            and zero if not hit.
+        n_pix_submap (int):  The number of pixels in a submap.
+        nside (int):  The Healpix NSIDE of the pixelization.
+        nest (bool):  If true, use NESTED ordering, else use RING.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return pixels_healpix(
+        quat_index,
+        quats,
+        shared_flags,
+        shared_flag_mask,
+        pixel_index,
+        pixels,
+        intervals,
+        hit_submaps,
+        n_pix_submap,
+        nside,
+        nest,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
diff --git a/src/toast/ops/pixels_healpix/kernels_jax.py b/src/toast/ops/pixels_healpix/kernels_jax.py
new file mode 100644
index 000000000..21f15a54d
--- /dev/null
+++ b/src/toast/ops/pixels_healpix/kernels_jax.py
@@ -0,0 +1,235 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import ALL, INTERVALS_JAX, JaxIntervals
+from ...jax.math import healpix, qarray
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def pixels_healpix_inner(hpix, quats, nest):
+    """
+    Compute the healpix pixel indices for the detectors.
+
+    Args:
+        hpix (HPIX_JAX): Healpix projection object.
+        quats (array, float64): Detector quaternion (size 4).
+        nest (bool): If True, then use NESTED ordering, else RING.
+
+    Returns:
+        pixels (array, int64): The detector pixel indices to store the result.
+    """
+    # constants
+    zaxis = jnp.array([0.0, 0.0, 1.0])
+
+    # initialize dir
+    dir = qarray.rotate_one_one(quats, zaxis)
+
+    # pixel computation
+    (phi, region, z, rtz) = healpix.vec2zphi(dir)
+    if nest:
+        pixel = healpix.zphi2nest(hpix, phi, region, z, rtz)
+    else:
+        pixel = healpix.zphi2ring(hpix, phi, region, z, rtz)
+
+    return pixel
+
+
+# maps over samples and detectors
+# pixels_healpix_inner = jax_xmap(pixels_healpix_inner,
+#                                 in_axes=[[...], # hpix
+#                                          ['detectors','intervals','interval_size',...], # quats
+#                                          [...]], # nest
+#                                 out_axes=['detectors','intervals','interval_size'])
+# TODO xmap is commented out for now due to a [bug with static argnum](https://github.com/google/jax/issues/10741)
+pixels_healpix_inner = jax.vmap(
+    pixels_healpix_inner, in_axes=[None, 0, None], out_axes=0
+)  # loop on interval_size
+pixels_healpix_inner = jax.vmap(
+    pixels_healpix_inner, in_axes=[None, 0, None], out_axes=0
+)  # loop on intervals
+pixels_healpix_inner = jax.vmap(
+    pixels_healpix_inner, in_axes=[None, 0, None], out_axes=0
+)  # loop on detectors
+
+
+def pixels_healpix_interval(
+    quat_index,
+    quats,
+    flags,
+    flag_mask,
+    pixel_index,
+    pixels,
+    hit_submaps,
+    n_pix_submap,
+    nside,
+    nest,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+):
+    """
+    Process all the intervals as a block.
+
+    Args:
+        quat_index (array, int): size n_det
+        quats (array, float64): The flat-packed array of detector quaternions (size ???*n_samp*4).
+        flags (array, uint8): size n_samp (or you shouldn't use flags)
+        flag_mask (uint8): integer used to select flags (not necesarely a boolean)
+        pixel_index (array, int): size n_det
+        pixels (array, int64): The detector pixel indices to store the result (size n_det*n_samp).
+        hit_submaps (array, uint8): The pointing flags (size ???).
+        n_pix_submap (int):
+        nside (int): Used to build the healpix projection object.
+        nest (bool): If True, then use NESTED ordering, else RING.
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+
+    Returns:
+        (pixels, hit_submaps)
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"pixels_healpix: jit-compiling.")
+
+    # flag
+    n_samp = pixels.shape[1]
+    use_flags = (flag_mask != 0) and (flags.size == n_samp)
+
+    # used to extract interval slices
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+
+    # computes the pixels and submap
+    hpix = healpix.HPIX_JAX(nside)
+    quats_interval = JaxIntervals.get(
+        quats, (quat_index, intervals, ALL)
+    )  # quats[quat_index,intervals,:]
+    pixels_interval = pixels_healpix_inner(hpix, quats_interval, nest)
+    sub_map = jnp.ravel(
+        pixels_interval // n_pix_submap
+    )  # flattened to index into the 1D hit_submaps
+    previous_hit_submaps_unflattened = jnp.reshape(
+        hit_submaps[sub_map], newshape=pixels_interval.shape
+    )  # unflattened to apply a 2D mask
+
+    # applies the flags
+    if use_flags:
+        # we pad with 1 such that values out of the interval will be flagged
+        flags_interval = JaxIntervals.get(
+            flags, intervals, padding_value=1
+        )  # flags[intervals]
+        is_flagged = (flags_interval & flag_mask) != 0
+        pixels_interval = jnp.where(is_flagged, -1, pixels_interval)
+        #
+        new_hit_submap_unflattened = jnp.where(
+            is_flagged, previous_hit_submaps_unflattened, 1
+        )
+    else:
+        # masks the padded values in sub_map
+        new_hit_submap_unflattened = jnp.where(
+            intervals.mask, previous_hit_submaps_unflattened, 1
+        )
+    new_hit_submap = jnp.ravel(new_hit_submap_unflattened)
+
+    # updates results and returns
+    pixels = jnp.array(pixels)  # TODO
+    hit_submaps = jnp.array(hit_submaps)  # TODO
+    hit_submaps = hit_submaps.at[sub_map].set(new_hit_submap)
+    pixels = JaxIntervals.set(
+        pixels, (pixel_index, intervals), pixels_interval
+    )  # pixels[pixel_index,intervals] = pixels_interval
+    return pixels, hit_submaps
+
+
+# jit compiling
+pixels_healpix_interval = jax.jit(
+    pixels_healpix_interval,
+    static_argnames=[
+        "flag_mask",
+        "n_pix_submap",
+        "nside",
+        "nest",
+        "intervals_max_length",
+    ],
+    donate_argnums=[5, 6],
+)  # donates pixels and hit_submap
+
+
+@kernel(impl=ImplementationType.JAX, name="pixels_healpix")
+def pixels_healpix_jax(
+    quat_index,
+    quats,
+    flags,
+    flag_mask,
+    pixel_index,
+    pixels,
+    intervals,
+    hit_submaps,
+    n_pix_submap,
+    nside,
+    nest,
+    use_accel,
+):
+    """
+    Compute the healpix pixel indices for the detectors.
+
+    Args:
+        quat_index (array, int): size n_det
+        quats (array, float64): The flat-packed array of detector quaternions (size ???*n_samp*4).
+        flags (array, uint8): size n_samp (or you shouldn't use flags)
+        flag_mask (uint8): integer used to select flags (not necesarely a boolean)
+        pixel_index (array, int): size n_det
+        pixels (array, int64): The detector pixel indices to store the result (size n_det*n_samp).
+        intervals (array, float64): size n_view
+        hit_submaps (array, uint8): The pointing flags (size ???).
+        n_pix_submap (int):
+        nside (int): Used to build the healpix projection object.
+        nest (bool): If True, then use NESTED ordering, else RING.
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (results are stored in pixels and hit_submaps).
+    """
+    # prepares inputs
+    if intervals.size == 0:
+        return  # deals with a corner case in tests
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    quat_index_input = MutableJaxArray.to_array(quat_index)
+    quats_input = MutableJaxArray.to_array(quats)
+    flags_input = MutableJaxArray.to_array(flags)
+    pixel_index_input = MutableJaxArray.to_array(pixel_index)
+    pixels_input = MutableJaxArray.to_array(pixels)
+    hit_submaps_input = MutableJaxArray.to_array(hit_submaps)
+
+    # runs computation
+    new_pixels, new_hit_submaps = pixels_healpix_interval(
+        quat_index_input,
+        quats_input,
+        flags_input,
+        flag_mask,
+        pixel_index_input,
+        pixels_input,
+        hit_submaps_input,
+        n_pix_submap,
+        nside,
+        nest,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+    )
+
+    # modifies output buffers in place
+    pixels[:] = new_pixels
+    hit_submaps[:] = new_hit_submaps
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_pointing_healpix", "ops_sim_ground", "ops_sim_satellite", "ops_demodulate");'
diff --git a/src/toast/ops/pixels_healpix/kernels_numpy.py b/src/toast/ops/pixels_healpix/kernels_numpy.py
new file mode 100644
index 000000000..5b7351a5a
--- /dev/null
+++ b/src/toast/ops/pixels_healpix/kernels_numpy.py
@@ -0,0 +1,45 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import healpy as hp
+import numpy as np
+
+from ... import qarray as qa
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="pixels_healpix")
+def pixels_healpix_numpy(
+    quat_index,
+    quats,
+    shared_flags,
+    shared_flag_mask,
+    pixel_index,
+    pixels,
+    intervals,
+    hit_submaps,
+    n_pix_submap,
+    nside,
+    nest,
+    use_accel=False,
+):
+    zaxis = np.array([0, 0, 1], dtype=np.float64)
+    for idet in range(len(quat_index)):
+        qidx = quat_index[idet]
+        pidx = pixel_index[idet]
+        for vw in intervals:
+            samples = slice(vw.first, vw.last + 1, 1)
+            dir = qa.rotate(quats[qidx][samples], zaxis)
+            pixels[pidx][samples] = hp.vec2pix(
+                nside,
+                dir[:, 0],
+                dir[:, 1],
+                dir[:, 2],
+                nest,
+            )
+            good = (shared_flags[samples] & shared_flag_mask) == 0
+            bad = np.logical_not(good)
+            sub_maps = pixels[pidx][samples][good] // n_pix_submap
+            hit_submaps[sub_maps] = 1
+            pixels[pidx][samples][bad] = -1
diff --git a/src/toast/ops/pixels_healpix.py b/src/toast/ops/pixels_healpix/pixels_healpix.py
similarity index 75%
rename from src/toast/ops/pixels_healpix.py
rename to src/toast/ops/pixels_healpix/pixels_healpix.py
index cc073e4ee..dd7ed667c 100644
--- a/src/toast/ops/pixels_healpix.py
+++ b/src/toast/ops/pixels_healpix/pixels_healpix.py
@@ -5,19 +5,15 @@
 import numpy as np
 import traitlets
 
-from .. import qarray as qa
-
-# This is just wrong- will be confusing until cleanup of
-# libtoast / _libtoast...
-from .._libtoast import healpix_pixels, pixels_healpix
-from ..healpix import Pixels
-from ..observation import default_values as defaults
-from ..pixels import PixelDistribution
-from ..timing import function_timer
-from ..traits import Bool, Instance, Int, Unicode, trait_docs
-from ..utils import Environment, Logger
-from .delete import Delete
-from .operator import Operator
+from ...accelerator import ImplementationType
+from ...healpix import Pixels
+from ...observation import default_values as defaults
+from ...pixels import PixelDistribution
+from ...timing import function_timer
+from ...traits import Bool, Instance, Int, Unicode, UseEnum, trait_docs
+from ...utils import Environment, Logger
+from ..operator import Operator
+from .kernels import pixels_healpix
 
 
 @trait_docs
@@ -71,8 +67,6 @@ class PixelsHealpix(Operator):
 
     single_precision = Bool(False, help="If True, use 32bit int in output")
 
-    use_python = Bool(False, help="If True, use python implementation")
-
     @traitlets.validate("detector_pointing")
     def _check_detector_pointing(self, proposal):
         detpointing = proposal["value"]
@@ -153,12 +147,12 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         env = Environment.get()
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.detector_pointing is None:
             raise RuntimeError("The detector_pointing trait must be set")
 
-        if self.use_python and use_accel:
-            raise RuntimeError("Cannot use accelerator with pure python implementation")
-
         if self._local_submaps is None and self.create_dist is not None:
             self._local_submaps = np.zeros(self._n_submap, dtype=np.uint8)
 
@@ -205,7 +199,6 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                         raise RuntimeError(msg)
 
             # Create (or re-use) output data for the pixels.
-
             if self.single_precision:
                 exists = ob.detdata.ensure(
                     self.pixels,
@@ -259,36 +252,27 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
 
             # Get the flags if needed.  Use the same flags as
             # detector pointing.  Use empty array in absence of flags
-            flags = np.array([], dtype=np.uint8)
-            if self.detector_pointing.shared_flags is not None:
-                flags = np.array(ob.shared[self.detector_pointing.shared_flags])
-
-            if self.use_python:
-                self._py_pixels_healpix(
-                    quat_indx,
-                    ob.detdata[quats_name].data,
-                    flags,
-                    self.detector_pointing.shared_flag_mask,
-                    pix_indx,
-                    ob.detdata[self.pixels].data,
-                    ob.intervals[self.view].data,
-                    hit_submaps,
-                )
+            if self.detector_pointing.shared_flags is None:
+                flags = np.array([], dtype=np.uint8)
             else:
-                pixels_healpix(
-                    quat_indx,
-                    ob.detdata[quats_name].data,
-                    flags,
-                    self.detector_pointing.shared_flag_mask,
-                    pix_indx,
-                    ob.detdata[self.pixels].data,
-                    ob.intervals[self.view].data,
-                    hit_submaps,
-                    self._n_pix_submap,
-                    self.nside,
-                    self.nest,
-                    use_accel,
-                )
+                flags = ob.shared[self.detector_pointing.shared_flags].data
+
+            pixels_healpix(
+                quat_indx,
+                ob.detdata[quats_name].data,
+                flags,
+                self.detector_pointing.shared_flag_mask,
+                pix_indx,
+                ob.detdata[self.pixels].data,
+                ob.intervals[self.view].data,
+                hit_submaps,
+                self._n_pix_submap,
+                self.nside,
+                self.nest,
+                impl=implementation,
+                use_accel=use_accel,
+            )
+
             if self._local_submaps is not None:
                 self._local_submaps[:] |= hit_submaps
 
@@ -338,53 +322,22 @@ def _requires(self):
         return req
 
     def _provides(self):
+        # FIXME quats generated are keyed in self.detector_pointing.quats (and not self.quats)
         prov = {
-            "detdata": [self.pixels],
+            "detdata": [self.pixels, self.detector_pointing.quats],
             "global": list(),
         }
         if self.create_dist is not None:
             prov["global"].append(self.create_dist)
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
     def _supports_accel(self):
-        return self.detector_pointing.supports_accel()
-
-    def _py_pixels_healpix(
-        self,
-        quat_indx,
-        quat_data,
-        flag_data,
-        flag_mask,
-        pix_indx,
-        pix_data,
-        intr_data,
-        hit_submaps,
-    ):
-        """Internal python implementation for comparison tests."""
-        zaxis = np.array([0, 0, 1], dtype=np.float64)
-        if self.nest:
-            for idet in range(len(quat_indx)):
-                qidx = quat_indx[idet]
-                pidx = pix_indx[idet]
-                for vw in intr_data:
-                    samples = slice(vw.first, vw.last + 1, 1)
-                    dir = qa.rotate(quat_data[qidx][samples], zaxis)
-                    pix_data[pidx][samples] = self.hpix.vec2nest(dir)
-                    good = (flag_data[samples] & flag_mask) == 0
-                    bad = np.logical_not(good)
-                    sub_maps = pix_data[pidx][samples][good] // self._n_pix_submap
-                    hit_submaps[sub_maps] = 1
-                    pix_data[pidx][samples][bad] = -1
-        else:
-            for idet in range(len(quat_indx)):
-                qidx = quat_indx[idet]
-                pidx = pix_indx[idet]
-                for vw in intr_data:
-                    samples = slice(vw.first, vw.last + 1, 1)
-                    dir = qa.rotate(quat_data[qidx][samples], zaxis)
-                    pix_data[pidx][samples] = self.hpix.vec2ring(dir)
-                    good = (flag_data[samples] & flag_mask) == 0
-                    bad = np.logical_not(good)
-                    sub_maps = pix_data[pidx][samples][good] // self._n_pix_submap
-                    hit_submaps[sub_maps] = 1
-                    pix_data[pidx][samples][bad] = -1
+        return True
diff --git a/src/toast/ops/pointing_detector/CMakeLists.txt b/src/toast/ops/pointing_detector/CMakeLists.txt
new file mode 100644
index 000000000..09e4fc5ac
--- /dev/null
+++ b/src/toast/ops/pointing_detector/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    pointing_detector.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/pointing_detector
+)
diff --git a/src/toast/ops/pointing_detector/__init__.py b/src/toast/ops/pointing_detector/__init__.py
new file mode 100644
index 000000000..f352cda29
--- /dev/null
+++ b/src/toast/ops/pointing_detector/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .pointing_detector import PointingDetectorSimple
diff --git a/src/toast/ops/pointing_detector/kernels.py b/src/toast/ops/pointing_detector/kernels.py
new file mode 100644
index 000000000..92eb7d80a
--- /dev/null
+++ b/src/toast/ops/pointing_detector/kernels.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ... import qarray as qa
+from ..._libtoast import pointing_detector as libtoast_pointing_detector
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import pointing_detector_numpy
+
+if use_accel_jax:
+    from .kernels_jax import pointing_detector_jax
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def pointing_detector(
+    focalplane,
+    boresight,
+    quat_index,
+    quats,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel=False,
+):
+    """Kernel for computing detector quaternion pointing.
+
+    Args:
+        focalplane (array):  Detector offset quaternions for each detector.
+        boresight (array):  Boresight quaternion pointing for each sample.
+        quat_index (array):  The index into the detector quaternion array for each
+            detector in the focalplane array.
+        quats (array):  The array of detector quaternions for each sample.
+        intervals (array):  The array of sample intervals.
+        shared_flags (array):  The array of common flags for each sample.
+        shared_flag_mask (int):  The flag mask to apply.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return pointing_detector(
+        focalplane,
+        boresight,
+        quat_index,
+        quats,
+        intervals,
+        shared_flags,
+        shared_flag_mask,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.COMPILED, name="pointing_detector")
+def pointing_detector_compiled(*args, use_accel=False):
+    return libtoast_pointing_detector(*args, use_accel)
diff --git a/src/toast/ops/pointing_detector/kernels_jax.py b/src/toast/ops/pointing_detector/kernels_jax.py
new file mode 100644
index 000000000..535662cce
--- /dev/null
+++ b/src/toast/ops/pointing_detector/kernels_jax.py
@@ -0,0 +1,171 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+from jax.experimental.maps import xmap as jax_xmap
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import ALL, INTERVALS_JAX, JaxIntervals
+from ...jax.math import qarray
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def pointing_detector_inner(focalplane, boresight, flag, mask):
+    """
+    Process a single detector and a single sample inside an interval.
+
+    Args:
+        focalplane (array, double): size 4
+        boresight (array, double): size 4
+        flag (uint8)
+        mask (uint8)
+
+    Returns:
+        quats (array, double): size 4
+    """
+    quats = jnp.where(
+        (flag & mask) == 0, qarray.mult_one_one(boresight, focalplane), focalplane
+    )
+    return quats
+
+
+# maps over intervals and detectors
+# pointing_detector_inner = jax_xmap(
+#    pointing_detector_inner,
+#    in_axes=[
+#        ["detectors", ...],  # focalplane
+#        ["intervals", "interval_size", ...],  # boresight
+#        ["intervals", "interval_size"],  # flags
+#        [...],
+#    ],  # mask
+#    out_axes=["detectors", "intervals", "interval_size", ...],
+# )
+# using vmap as the static arguments triggers the following error:
+# "ShardingContext cannot be used with xmap"
+# TODO revisit once this issue is solved [bug with static argnum](https://github.com/google/jax/issues/10741)
+pointing_detector_inner = jax.vmap(
+    pointing_detector_inner, in_axes=[None, 0, 0, None], out_axes=0
+)  # interval_size
+pointing_detector_inner = jax.vmap(
+    pointing_detector_inner, in_axes=[None, 0, 0, None], out_axes=0
+)  # intervals
+pointing_detector_inner = jax.vmap(
+    pointing_detector_inner, in_axes=[0, None, None, None], out_axes=0
+)  # detectors
+
+
+def pointing_detector_interval(
+    focalplane,
+    boresight,
+    quat_index,
+    quats,
+    shared_flags,
+    shared_flag_mask,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+):
+    """
+    Process all the intervals as a block.
+
+    Args:
+        focalplane (array, double): size n_det*4
+        boresight (array, double): size n_samp*4
+        quat_index (array, int): size n_det
+        quats (array, double): size n_det*n_samp*4
+        shared_flags (array, uint8): size n_samp
+        shared_flag_mask (uint8)
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+
+    Returns:
+        quats
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"pointing_detector: jit-compiling.")
+
+    # extract interval slices
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    boresight_interval = JaxIntervals.get(
+        boresight, (intervals, ALL)
+    )  # boresight[intervals,:]
+    shared_flags_interval = JaxIntervals.get(
+        shared_flags, intervals
+    )  # shared_flags[intervals]
+
+    # process the interval then updates quats in place
+    new_quats_interval = pointing_detector_inner(
+        focalplane, boresight_interval, shared_flags_interval, shared_flag_mask
+    )
+    quats = JaxIntervals.set(
+        quats, (quat_index, intervals, ALL), new_quats_interval
+    )  # quats[quat_index,intervals,:] = new_quats_interval
+    return quats
+
+
+# jit compiling
+pointing_detector_interval = jax.jit(
+    pointing_detector_interval,
+    static_argnames=["shared_flag_mask", "intervals_max_length"],
+    donate_argnums=[3],
+)  # donates quats
+
+
+@kernel(impl=ImplementationType.JAX, name="pointing_detector")
+def pointing_detector_jax(
+    focalplane,
+    boresight,
+    quat_index,
+    quats,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel=False,
+):
+    """
+    Args:
+        focalplane (array, double): size n_det*4
+        boresight (array, double): size n_samp*4
+        quat_index (array, int): size n_det
+        quats (array, double): size n_det*n_samp*4
+        intervals (array, Interval): The intervals to modify (size n_view)
+        shared_flags (array, uint8): size n_samp
+        shared_flag_mask (uint8)
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in quats).
+    """
+    # prepares inputs
+    if intervals.size == 0:
+        return  # deals with a corner case in tests
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    focalplane_input = MutableJaxArray.to_array(focalplane)
+    boresight_input = MutableJaxArray.to_array(boresight)
+    quat_index_input = MutableJaxArray.to_array(quat_index)
+    quats_input = MutableJaxArray.to_array(quats)
+    shared_flags_input = MutableJaxArray.to_array(shared_flags)
+
+    # runs computation
+    quats[:] = pointing_detector_interval(
+        focalplane_input,
+        boresight_input,
+        quat_index_input,
+        quats_input,
+        shared_flags_input,
+        shared_flag_mask,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+    )
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_pointing_healpix", "ops_demodulate")'
diff --git a/src/toast/ops/pointing_detector/kernels_numpy.py b/src/toast/ops/pointing_detector/kernels_numpy.py
new file mode 100644
index 000000000..60f281d45
--- /dev/null
+++ b/src/toast/ops/pointing_detector/kernels_numpy.py
@@ -0,0 +1,30 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ... import qarray as qa
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="pointing_detector")
+def pointing_detector_numpy(
+    focalplane,
+    boresight,
+    quat_index,
+    quats,
+    intervals,
+    shared_flags,
+    shared_flag_mask,
+    use_accel=False,
+):
+    for idet in range(len(quat_index)):
+        qidx = quat_index[idet]
+        for vw in intervals:
+            samples = slice(vw.first, vw.last + 1, 1)
+            bore = np.array(boresight[samples])
+            if shared_flags is not None:
+                good = (shared_flags[samples] & shared_flag_mask) == 0
+                bore[np.invert(good)] = np.array([0, 0, 0, 1], dtype=np.float64)
+            quats[qidx][samples] = qa.mult(bore, focalplane[idet])
diff --git a/src/toast/ops/pointing_detector.py b/src/toast/ops/pointing_detector/pointing_detector.py
similarity index 74%
rename from src/toast/ops/pointing_detector.py
rename to src/toast/ops/pointing_detector/pointing_detector.py
index 1ce3af531..52f637a2f 100644
--- a/src/toast/ops/pointing_detector.py
+++ b/src/toast/ops/pointing_detector/pointing_detector.py
@@ -5,13 +5,14 @@
 import numpy as np
 import traitlets
 
-from .. import qarray as qa
-from .._libtoast import pointing_detector
-from ..observation import default_values as defaults
-from ..timing import function_timer
-from ..traits import Bool, Int, Unicode, trait_docs
-from ..utils import Environment, Logger
-from .operator import Operator
+from ... import qarray as qa
+from ...accelerator import ImplementationType
+from ...observation import default_values as defaults
+from ...timing import function_timer
+from ...traits import Bool, Int, Unicode, UseEnum, trait_docs
+from ...utils import Logger
+from ..operator import Operator
+from .kernels import pointing_detector
 
 
 @trait_docs
@@ -58,8 +59,6 @@ class PointingDetectorSimple(Operator):
         help="The output coordinate system ('C', 'E', 'G')",
     )
 
-    use_python = Bool(False, help="If True, use python implementation")
-
     @traitlets.validate("coord_in")
     def _check_coord_in(self, proposal):
         check = proposal["value"]
@@ -90,8 +89,8 @@ def __init__(self, **kwargs):
     def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
-        if self.use_python and use_accel:
-            raise RuntimeError("Cannot use accelerator with pure python implementation")
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
 
         coord_rot = None
         if self.coord_in is None:
@@ -142,8 +141,7 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                     log.verbose(msg)
                 continue
 
-            # FIXME:  temporary hack until instrument classes are also pre-staged
-            # to GPU
+            # FIXME:  temporary hack until instrument classes are also pre-staged to GPU
             focalplane = ob.telescope.focalplane
             fp_quats = np.zeros((len(dets), 4), dtype=np.float64)
             for idet, d in enumerate(dets):
@@ -167,26 +165,17 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
 
             # FIXME: handle coordinate transforms here too...
 
-            if self.use_python:
-                self._py_pointing_detector(
-                    fp_quats,
-                    ob.shared[self.boresight].data,
-                    quat_indx,
-                    ob.detdata[self.quats].data,
-                    ob.intervals[self.view].data,
-                    flags,
-                )
-            else:
-                pointing_detector(
-                    fp_quats,
-                    ob.shared[self.boresight].data,
-                    quat_indx,
-                    ob.detdata[self.quats].data,
-                    ob.intervals[self.view].data,
-                    flags,
-                    self.shared_flag_mask,
-                    use_accel,
-                )
+            pointing_detector(
+                fp_quats,
+                ob.shared[self.boresight].data,
+                quat_indx,
+                ob.detdata[self.quats].data,
+                ob.intervals[self.view].data,
+                flags,
+                self.shared_flag_mask,
+                impl=implementation,
+                use_accel=use_accel,
+            )
 
         return
 
@@ -214,25 +203,16 @@ def _provides(self):
         }
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
     def _supports_accel(self):
         return True
 
-    def _py_pointing_detector(
-        self,
-        fp_quats,
-        bore_data,
-        quat_indx,
-        quat_data,
-        intr_data,
-        flag_data,
-    ):
-        """Internal python implementation for comparison tests."""
-        for idet in range(len(quat_indx)):
-            qidx = quat_indx[idet]
-            for vw in intr_data:
-                samples = slice(vw.first, vw.last + 1, 1)
-                bore = np.array(bore_data[samples])
-                if self.shared_flags is not None:
-                    good = (flag_data[samples] & self.shared_flag_mask) == 0
-                    bore[np.invert(good)] = np.array([0, 0, 0, 1], dtype=np.float64)
-                quat_data[qidx][samples] = qa.mult(bore, fp_quats[idet])
+    def _supports_accel(self):
+        return True
diff --git a/src/toast/ops/polyfilter/CMakeLists.txt b/src/toast/ops/polyfilter/CMakeLists.txt
new file mode 100644
index 000000000..12409f475
--- /dev/null
+++ b/src/toast/ops/polyfilter/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    polyfilter.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/polyfilter
+)
diff --git a/src/toast/ops/polyfilter/__init__.py b/src/toast/ops/polyfilter/__init__.py
new file mode 100644
index 000000000..09d58f17e
--- /dev/null
+++ b/src/toast/ops/polyfilter/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .polyfilter import CommonModeFilter, PolyFilter, PolyFilter2D
diff --git a/src/toast/ops/polyfilter/kernels.py b/src/toast/ops/polyfilter/kernels.py
new file mode 100644
index 000000000..9c0b8fb17
--- /dev/null
+++ b/src/toast/ops/polyfilter/kernels.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ..._libtoast import filter_poly2D as libtoast_filter_poly2D
+from ..._libtoast import filter_polynomial as libtoast_filter_polynomial
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import filter_poly2D_numpy, filter_polynomial_numpy
+
+if use_accel_jax:
+    from .kernels_jax import filter_poly2D_jax, filter_polynomial_jax
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def filter_polynomial(
+    order,
+    flags,
+    signals_list,
+    starts,
+    stops,
+    use_accel=False,
+):
+    """Kernel to fit and subtract a polynomial from one or more signals.
+
+    Args:
+        order (int):  The order of the polynomial.
+        flags (numpy array, uint8):  The common flags to use for all signals
+        signals_list (list of numpy array of double):  A list of float64 arrays containing the signals.
+        starts (numpy array, int64):  The start samples of each scan.
+        stops (numpy array, int64):  The stop samples of each scan.
+
+    Returns:
+        None
+
+    """
+    return filter_polynomial(
+        order,
+        flags,
+        signals_list,
+        starts,
+        stops,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def filter_poly2D(
+    det_groups,
+    templates,
+    signals,
+    masks,
+    coeff,
+    use_accel=False,
+):
+    """Kernel for solving 2D polynomial coefficients at each sample.
+
+    Args:
+        det_groups (numpy array, int32):  The group index for each of the N_detectors detector index.
+        templates (numpy array, float64):  The N_detectors x N_modes templates.
+        signals (numpy array, float64):  The N_sample x N_detector data.
+        masks (numpy array, uint8):  The N_sample x N_detector mask.
+        coeff (numpy array, float64):  The N_sample x N_group x N_mode output coefficients.
+
+    Returns:
+        None
+
+    """
+    return filter_poly2D(
+        det_groups,
+        templates,
+        signals,
+        masks,
+        coeff,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.COMPILED, name="filter_polynomial")
+def filter_polynomial_compiled(*args, use_accel=False):
+    return libtoast_filter_polynomial(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="filter_poly2D")
+def filter_poly2D_compiled(*args, use_accel=False):
+    return libtoast_filter_poly2D(*args, use_accel)
diff --git a/src/toast/ops/polyfilter/kernels_jax.py b/src/toast/ops/polyfilter/kernels_jax.py
new file mode 100644
index 000000000..51cc1851e
--- /dev/null
+++ b/src/toast/ops/polyfilter/kernels_jax.py
@@ -0,0 +1,218 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax.experimental.maps import xmap as jax_xmap
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.mutableArray import MutableJaxArray
+
+
+def filter_poly2D_sample_group(
+    igroup, det_groups, templates, signals_sample, masks_sample
+):
+    """
+    filter_poly2D for a given group and sample
+
+    Args:
+        igroup (uint): index of the group
+        det_groups (numpy array, int32):  The group index for each of the N_detectors detector index.
+        templates (numpy array, float64):  The N_detectors x N_modes templates.
+        signals_samples (numpy array, float64):  The N_detector data.
+        masks_sample (numpy array, uint8):  The N_detector mask.
+
+    Returns:
+        coeff (numpy array, float64):  The N_mode output coefficients.
+    """
+    # Masks detectors not in this group
+    masks_group = jnp.where(det_groups != igroup, 0.0, masks_sample)
+
+    # rhs = (mask * templates).T  @  (mask * signals.T) = (mask * templates).T  @  signals.T
+    # A = (mask * templates).T  @  (mask * templates)
+    masked_template = masks_group[:, jnp.newaxis] * templates  # N_detectors x N_modes
+    rhs = jnp.dot(masked_template.T, signals_sample.T)  # N_modes
+    A = jnp.dot(masked_template.T, masked_template)  # N_modes x N_modes
+
+    # Fits the coefficients
+    (coeff_sample_group, _residue, _rank, _singular_values) = jnp.linalg.lstsq(
+        A, rhs, rcond=1e-3
+    )
+    # Sometimes the mask will be all zeroes in which case A=0 and rhs=0 causing the coeffs to be nan
+    # We thus replace nans with 0s (Numpy does it by default)
+    coeff_sample_group = jnp.nan_to_num(coeff_sample_group, nan=0.0)
+
+    return coeff_sample_group
+
+
+def filter_poly2D_coeffs(ngroup, det_groups, templates, signals, masks):
+    """
+    Args:
+        ngroup (uint): number of groups
+        det_groups (numpy array, int32):  The group index for each of the N_detectors detector index.
+        templates (numpy array, float64):  The N_detectors x N_modes templates.
+        signals (numpy array, float64):  The N_sample x N_detector data.
+        masks (numpy array, uint8):  The N_sample x N_detector mask.
+
+    Returns:
+        coeff (numpy array, float64):  The N_sample x N_group x N_mode output coefficients.
+    """
+    # batch on group and sample dimenssions
+    filter_poly2D_sample_group_batched = jax_xmap(
+        filter_poly2D_sample_group,
+        in_axes=[
+            ["group"],  # igroup
+            [...],  # det_groups
+            [...],  # templates
+            ["sample", ...],  # signals
+            ["sample", ...],
+        ],  # masks
+        out_axes=["sample", "group", ...],
+    )
+
+    # runs for all the groups and samples simultaneously
+    igroup = jnp.arange(start=0, stop=ngroup)
+    return filter_poly2D_sample_group_batched(
+        igroup, det_groups, templates, signals, masks
+    )
+
+
+# JIT compiles the code
+filter_poly2D_coeffs = jax.jit(filter_poly2D_coeffs, static_argnames="ngroup")
+
+
+@kernel(impl=ImplementationType.JAX, name="filter_poly2D")
+def filter_poly2D_jax(det_groups, templates, signals, masks, coeff):
+    """
+    Solves for 2D polynomial coefficients at each sample.
+
+    Args:
+        det_groups (numpy array, int32):  The group index for each of the N_detectors detector index.
+        templates (numpy array, float64):  The N_detectors x N_modes templates.
+        signals (numpy array, float64):  The N_sample x N_detector data.
+        masks (numpy array, uint8):  The N_sample x N_detector mask.
+        coeff (numpy array, float64):  The N_sample x N_group x N_mode output coefficients.
+
+    Returns:
+        None: The coefficients are updated in place.
+    """
+    # does computation with JAX function and assign result to coeffs
+    ngroup = coeff.shape[1]
+    det_groups_input = MutableJaxArray.to_array(det_groups)
+    templates_input = MutableJaxArray.to_array(templates)
+    signals_input = MutableJaxArray.to_array(signals)
+    masks_input = MutableJaxArray.to_array(masks)
+    coeff[:] = filter_poly2D_coeffs(
+        ngroup, det_groups_input, templates_input, signals_input, masks_input
+    )
+
+
+def filter_polynomial_interval(flags_interval, signals_interval, order):
+    """
+    Process a single interval
+    Return signals_interval modified by applying the polynomial filter
+    """
+    # problem size
+    norder = order + 1
+    scanlen = flags_interval.size
+
+    # Build the full template matrix used to clean the signal.
+    # We subtract the template value even from flagged samples to support point source masking etc.
+    full_templates = jnp.empty(shape=(scanlen, norder))  # scanlen*norder
+    # deals with order 0
+    if norder > 0:
+        # full_templates[:,0] = 1
+        full_templates = full_templates.at[:, 0].set(1)
+    # deals with order 1
+    if norder > 1:
+        # defines x
+        indices = jnp.arange(start=0, stop=scanlen)
+        xstart = (1.0 / scanlen) - 1.0
+        dx = 2.0 / scanlen
+        x = xstart + dx * indices
+        # full_templates[:,1] = x
+        full_templates = full_templates.at[:, 1].set(x)
+    # deals with other orders
+    # this loop will be unrolled but this should be okay as `order` is likely small
+    for iorder in range(2, norder):
+        previous_previous_order = full_templates[:, iorder - 2]
+        previous_order = full_templates[:, iorder - 1]
+        current_order = (
+            (2 * iorder - 1) * x * previous_order
+            - (iorder - 1) * previous_previous_order
+        ) / iorder
+        # full_templates[:,iorder] = current_order
+        full_templates = full_templates.at[:, iorder].set(current_order)
+
+    # Assemble the flagged template matrix used in the linear regression
+    # by zeroing out the rows that are flagged
+    valid_rows = flags_interval == 0
+    masked_templates = jnp.where(
+        valid_rows[:, jnp.newaxis], full_templates, 0.0
+    )  # nb_zero_flags*norder
+
+    # Square the template matrix for A^T.A
+    invcov = jnp.dot(masked_templates.T, masked_templates)  # norder*norder
+
+    # Project the signals against the templates
+    # we do not mask flagged signals as they are multiplied with zeros anyway
+    proj = jnp.dot(masked_templates.T, signals_interval)  # norder*nsignal
+
+    # Fit the templates against the data
+    (x, _residue, _rank, _singular_values) = jnp.linalg.lstsq(
+        invcov, proj, rcond=1e-3
+    )  # norder*nsignal
+
+    # computes the value to be subtracted from the signals
+    return signals_interval - jnp.dot(full_templates, x)
+
+
+# JIT compiles the code
+filter_polynomial_interval = jax.jit(
+    filter_polynomial_interval, static_argnames=["order"]
+)
+
+
+@kernel(impl=ImplementationType.JAX, name="filter_polynomial")
+def filter_polynomial_jax(order, flags, signals_list, starts, stops):
+    """
+    Fit and subtract a polynomial from one or more signals.
+
+    Args:
+        order (int):  The order of the polynomial.
+        flags (numpy array, uint8):  The common flags to use for all signals
+        signals_list (list of numpy array of double):  A list of float64 arrays containing the signals.
+        starts (numpy array, int64):  The start samples of each scan.
+        stops (numpy array, int64):  The stop samples of each scan.
+
+    Returns:
+        None: The signals are updated in place.
+    """
+    # validate order
+    if order < 0:
+        return
+
+    # converts signals from a list into a numpy array to avoid having to loop over them
+    signals = np.array(signals_list).T  # n*nsignal
+
+    # loop over intervals, this is fine as long as there are only few intervals
+    # TODO port to JaxIntervals, could be done with vmap and setting padding of flags_interval to 1
+    for start, stop in zip(starts, stops):
+        # validates interval
+        start = np.maximum(0, start)
+        stop = np.minimum(flags.size - 1, stop) + 1
+        if stop <= start:
+            continue
+        # extracts the intervals from flags and signals
+        flags_interval = flags[start:stop]  # scanlen
+        signals_interval = signals[start:stop, :]  # scanlen*nsignal
+        # updates signal interval
+        signals[start:stop, :] = filter_polynomial_interval(
+            flags_interval, signals_interval, order
+        )
+
+    # puts new signals back into the list
+    for isignal, signal in enumerate(signals_list):
+        signal[:] = signals[:, isignal]
diff --git a/src/toast/ops/polyfilter/kernels_numpy.py b/src/toast/ops/polyfilter/kernels_numpy.py
new file mode 100644
index 000000000..9c756769c
--- /dev/null
+++ b/src/toast/ops/polyfilter/kernels_numpy.py
@@ -0,0 +1,101 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="filter_polynomial")
+def filter_polynomial_numpy(order, flags, signals_list, starts, stops, use_accel=False):
+    # validate order
+    if order < 0:
+        return
+
+    # problem size
+    n = flags.size
+    norder = order + 1
+
+    # converts signal into a numpy array to avoid having to loop over them
+    signals = np.array(signals_list).T  # n*nsignal
+
+    for start, stop in zip(starts, stops):
+        # validates interval
+        start = np.maximum(0, start)
+        stop = np.minimum(n - 1, stop)
+        if stop < start:
+            continue
+        scanlen = stop - start + 1
+
+        # extracts the signals that will be impacted by this interval
+        signals_interval = signals[start : (stop + 1), :]  # scanlen*nsignal
+
+        # set aside the indexes of the zero flags to be used as a mask
+        flags_interval = flags[start : (stop + 1)]  # scanlen
+        zero_flags = np.where(flags_interval == 0)
+        nb_zero_flags = zero_flags[0].size
+        if nb_zero_flags == 0:
+            continue
+
+        # Build the full template matrix used to clean the signal.
+        # We subtract the template value even from flagged samples to
+        # support point source masking etc.
+        # NOTE: we could recycle full_template if the previous scanlen is identical
+        full_templates = np.empty(shape=(scanlen, norder))  # scanlen*norder
+        # deals with order 0
+        if norder > 0:
+            full_templates[:, 0] = 1
+        # deals with order 1
+        if norder > 1:
+            # defines x
+            xstart = (1.0 / scanlen) - 1.0
+            xstop = (1.0 / scanlen) + 1.0
+            dx = 2.0 / scanlen
+            x = np.arange(start=xstart, stop=xstop, step=dx)
+            # sets ordre 1
+            full_templates[:, 1] = x
+        # deals with other orders
+        # NOTE: this formulation is inherently sequential but this should be okay as `order` is likely small
+        for iorder in range(2, norder):
+            previous_previous_order = full_templates[:, iorder - 2]
+            previous_order = full_templates[:, iorder - 1]
+            full_templates[:, iorder] = (
+                (2 * iorder - 1) * x * previous_order
+                - (iorder - 1) * previous_previous_order
+            ) / iorder
+
+        # Assemble the flagged template matrix used in the linear regression
+        masked_templates = full_templates[zero_flags]  # nb_zero_flags*norder
+
+        # Square the template matrix for A^T.A
+        invcov = np.dot(masked_templates.T, masked_templates)  # norder*norder
+
+        # Project the signals against the templates
+        masked_signals = signals_interval[zero_flags]  # nb_zero_flags*nsignal
+        proj = np.dot(masked_templates.T, masked_signals)  # norder*nsignal
+
+        # Fit the templates against the data
+        # by minimizing the norm2 of the difference and the solution vector
+        (x, _residue, _rank, _singular_values) = np.linalg.lstsq(
+            invcov, proj, rcond=1e-3
+        )  # norder*nsignal
+        signals_interval -= np.dot(full_templates, x)
+
+    # puts resulting signals back into list form
+    for isignal, signal in enumerate(signals_list):
+        signal[:] = signals[:, isignal]
+
+
+@kernel(impl=ImplementationType.NUMPY, name="filter_poly2D")
+def filter_poly2D_numpy(det_groups, templates, signals, masks, coeff, use_accel=False):
+    ngroup = coeff.shape[1]
+    nsample = signals.shape[0]
+    for isample in range(nsample):
+        for igroup in range(ngroup):
+            good = det_groups == igroup
+            mask = masks[isample, good]
+            t = templates[good].T.copy() * mask
+            proj = np.dot(t, signals[isample, good] * mask)
+            ccinv = np.dot(t, t.T)
+            coeff[isample, igroup] = np.linalg.lstsq(ccinv, proj, rcond=1.0e-6)[0]
diff --git a/src/toast/ops/polyfilter.py b/src/toast/ops/polyfilter/polyfilter.py
similarity index 93%
rename from src/toast/ops/polyfilter.py
rename to src/toast/ops/polyfilter/polyfilter.py
index 2b1e5cf3d..07091c6c5 100644
--- a/src/toast/ops/polyfilter.py
+++ b/src/toast/ops/polyfilter/polyfilter.py
@@ -11,13 +11,14 @@
 import traitlets
 from astropy import units as u
 
-from .. import qarray as qa
-from .._libtoast import filter_poly2D, filter_polynomial, subtract_mean, sum_detectors
-from ..mpi import MPI, Comm, MPI_Comm, use_mpi
-from ..observation import default_values as defaults
-from ..timing import function_timer
-from ..traits import Bool, Int, Unicode, trait_docs
-from ..utils import (
+from ... import qarray as qa
+from ..._libtoast import subtract_mean, sum_detectors
+from ...accelerator import ImplementationType
+from ...mpi import MPI, Comm, MPI_Comm, use_mpi
+from ...observation import default_values as defaults
+from ...timing import function_timer
+from ...traits import Bool, Dict, Instance, Int, Quantity, Unicode, UseEnum, trait_docs
+from ...utils import (
     AlignedF64,
     AlignedU8,
     Environment,
@@ -26,7 +27,8 @@
     Timer,
     dtype_to_aligned,
 )
-from .operator import Operator
+from ..operator import Operator
+from .kernels import filter_poly2D, filter_polynomial
 
 XAXIS, YAXIS, ZAXIS = np.eye(3)
 
@@ -87,10 +89,6 @@ class PolyFilter2D(Operator):
         None, allow_none=True, help="Which focalplane key to match"
     )
 
-    use_python = Bool(
-        False, help="If True, use a pure python implementation for testing."
-    )
-
     @traitlets.validate("shared_flag_mask")
     def _check_shared_flag_mask(self, proposal):
         check = proposal["value"]
@@ -110,9 +108,12 @@ def __init__(self, **kwargs):
         return
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         gt = GlobalTimers.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if detectors is not None:
             raise RuntimeError("PolyFilter2D cannot be run on subsets of detectors")
         norder = self.order + 1
@@ -323,23 +324,17 @@ def _exec(self, data, detectors=None, **kwargs):
 
                 gt.stop("Poly2D:  Accumulate templates")
 
-                if self.use_python:
-                    gt.start("Poly2D:  Solve templates (with python)")
-                    for isample in range(nsample):
-                        for group, igroup in group_ids.items():
-                            good = group_det == igroup
-                            mask = masks[isample, good]
-                            t = templates[good].T.copy() * mask
-                            proj = np.dot(t, signals[isample, good] * mask)
-                            ccinv = np.dot(t, t.T)
-                            coeff[isample, igroup] = np.linalg.lstsq(
-                                ccinv, proj, rcond=1.0e-6
-                            )[0]
-                    gt.stop("Poly2D:  Solve templates (with python)")
-                else:
-                    gt.start("Poly2D:  Solve templates")
-                    filter_poly2D(det_groups, templates, signals, masks, coeff)
-                    gt.stop("Poly2D:  Solve templates")
+                gt.start("Poly2D:  Solve templates")
+                filter_poly2D(
+                    det_groups,
+                    templates,
+                    signals,
+                    masks,
+                    coeff,
+                    impl=implementation,
+                    use_accel=use_accel,
+                )
+                gt.stop("Poly2D:  Solve templates")
 
                 gt.start("Poly2D:  Update detector flags")
 
@@ -494,9 +489,12 @@ def __init__(self, **kwargs):
         return
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.pattern is None:
             pat = None
         else:
@@ -553,14 +551,26 @@ def _exec(self, data, detectors=None, **kwargs):
                     signals.append(signal)
                 else:
                     filter_polynomial(
-                        self.order, last_flags, signals, local_starts, local_stops
+                        self.order,
+                        last_flags,
+                        signals,
+                        local_starts,
+                        local_stops,
+                        impl=implementation,
+                        use_accel=use_accel,
                     )
                     signals = [signal]
                 last_flags = flags.copy()
 
             if len(signals) > 0:
                 filter_polynomial(
-                    self.order, last_flags, signals, local_starts, local_stops
+                    self.order,
+                    last_flags,
+                    signals,
+                    local_starts,
+                    local_stops,
+                    impl=implementation,
+                    use_accel=use_accel,
                 )
 
             # Optionally flag unfiltered data
@@ -732,7 +742,7 @@ def _re_redistribute(self, data, obs, timer, log, temp_ob):
         return
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         """Apply the common mode filter to the signal.
 
         Args:
@@ -747,6 +757,9 @@ def _exec(self, data, detectors=None, **kwargs):
         timer.start()
         pat = re.compile(self.pattern)
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         for obs in data.obs:
             comm, temp_ob = self._redistribute(data, obs, timer, log)
 
diff --git a/src/toast/ops/scan_map/CMakeLists.txt b/src/toast/ops/scan_map/CMakeLists.txt
new file mode 100644
index 000000000..6826bbb66
--- /dev/null
+++ b/src/toast/ops/scan_map/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    scan_map.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/scan_map
+)
diff --git a/src/toast/ops/scan_map/__init__.py b/src/toast/ops/scan_map/__init__.py
new file mode 100644
index 000000000..96057d263
--- /dev/null
+++ b/src/toast/ops/scan_map/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .scan_map import ScanMap, ScanMask, ScanScale
diff --git a/src/toast/ops/scan_map/kernels.py b/src/toast/ops/scan_map/kernels.py
new file mode 100644
index 000000000..bdda9fee3
--- /dev/null
+++ b/src/toast/ops/scan_map/kernels.py
@@ -0,0 +1,172 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import healpy as hp
+import numpy as np
+
+from ... import qarray as qa
+from ..._libtoast import AlignedF64, AlignedI64
+from ..._libtoast import scan_map_float32 as libtoast_scan_map_float32
+from ..._libtoast import scan_map_float64 as libtoast_scan_map_float64
+from ..._libtoast import scan_map_int32 as libtoast_scan_map_int32
+from ..._libtoast import scan_map_int64 as libtoast_scan_map_int64
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import scan_map_numpy
+
+if use_accel_jax:
+    from .kernels_jax import scan_map_jax
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def scan_map(
+    global2local,
+    n_pix_submap,
+    mapdata,
+    det_data,
+    det_data_index,
+    pixels,
+    pixels_index,
+    weights,
+    weight_index,
+    intervals,
+    data_scale=1.0,
+    should_zero=False,
+    should_subtract=False,
+    should_scale=False,
+    use_accel=False,
+):
+    """Kernel for scanning a map into timestreams.
+
+    This uses a local piece of a distributed map and the pointing matrix for local
+    detectors to generate timestream values.
+
+    Args:
+        global2local (array):  The mapping from global submap to local submap index.
+        n_pix_submap (int):  The number of pixels per submap.
+        mapdata (array):  The local piece of the map.
+        det_data (array):  The detector data at each sample for each detector.
+        data_index (array):  The index into the data array for each detector.
+        pixels (array):  The array of detector pixels for each sample.
+        pixel_index (array):  The index into the detector pixel array for each
+            detector.
+        weights (array):  The array of I, Q, and U weights at each sample for each
+            detector.
+        weight_index (array):  The index into the weights array for each detector.
+        intervals (array):  The array of sample intervals.
+        data_scale (float):  Scale factor to apply to map data before accumulation.
+        should_zero (bool):  If True, set the detector data to zero before accumulation.
+        should_subtract (bool):  If True, subtract rather than add to detector data.
+        should_scale (bool):  If True, scale detector data by the scanned map values.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None.
+
+    """
+    return scan_map(
+        global2local,
+        n_pix_submap,
+        mapdata,
+        det_data,
+        det_data_index,
+        pixels,
+        pixels_index,
+        weights,
+        weight_index,
+        intervals,
+        data_scale=1.0,
+        should_zero=should_zero,
+        should_subtract=should_subtract,
+        should_scale=should_scale,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+# FIXME:  This "compiled" kernel will migrate fully to the _libtoast extension.
+
+
+@kernel(impl=ImplementationType.COMPILED, name="scan_map")
+def scan_map_compiled(
+    global2local,
+    n_pix_submap,
+    mapdata,
+    det_data,
+    det_data_index,
+    pixels,
+    pixels_index,
+    weights,
+    weight_index,
+    intervals,
+    data_scale=1.0,
+    should_zero=False,
+    should_subtract=False,
+    should_scale=False,
+    use_accel=False,
+):
+    if mapdata.dtype.char == "d":
+        fcomp = libtoast_scan_map_float64
+    elif mapdata.dtype.char == "f":
+        fcomp = libtoast_scan_map_float32
+    elif mapdata.dtype.char == "i":
+        fcomp = libtoast_scan_map_int32
+    elif mapdata.dtype.char == "l":
+        fcomp = libtoast_scan_map_int64
+    else:
+        msg = f"Compiled version of scan_map does not support array "
+        msg += f"type '{mapdata.dtype.char}'"
+        raise NotImplementedError(msg)
+
+    nmap = 1
+    if len(weights.shape) > 2:
+        nmap = weights.shape[2]
+    for idet in range(len(det_data_index)):
+        pix = pixels[pixels_index[idet]]
+        wts = weights[weight_index[idet]]
+        tod = det_data[det_data_index[idet]]
+        for view in intervals:
+            view_samples = view.last + 1 - view.first
+            vslice = slice(view.first, view.last + 1)
+
+            # Temp buffers
+            local_tod_raw = AlignedF64.zeros(view_samples)
+            local_tod = local_tod_raw.array()
+            local_submap_pix_raw = AlignedI64.zeros(view_samples)
+            local_submap_pix = local_submap_pix_raw.array()
+            local_submap_pix[:] = -1
+            local_submap_raw = AlignedI64.zeros(view_samples)
+            local_submap = local_submap_raw.array()
+            local_submap[:] = -1
+
+            # Get local submaps and pixel indices within each submap
+            good = pix[vslice] >= 0
+            local_submap_pix[good] = pix[vslice][good] % n_pix_submap
+            local_submap[good] = global2local[pix[vslice][good] // n_pix_submap]
+
+            fcomp(
+                n_pix_submap,
+                nmap,
+                local_submap,
+                local_submap_pix,
+                mapdata.reshape((-1,)),
+                wts[vslice].reshape((-1,)),
+                local_tod,
+            )
+
+            # Accumulate
+            if should_zero:
+                tod[vslice] = 0
+            if should_subtract:
+                tod[vslice][good] -= local_tod[good]
+            elif should_scale:
+                tod[vslice][good] *= local_tod[good]
+            else:
+                tod[vslice][good] += local_tod[good]
+
+            del local_submap
+            del local_submap_raw
+            del local_submap_pix
+            del local_submap_pix_raw
+            del local_tod
+            del local_tod_raw
diff --git a/src/toast/ops/scan_map/kernels_jax.py b/src/toast/ops/scan_map/kernels_jax.py
new file mode 100644
index 000000000..6ffac49bb
--- /dev/null
+++ b/src/toast/ops/scan_map/kernels_jax.py
@@ -0,0 +1,291 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import ALL, INTERVALS_JAX, JaxIntervals
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def global_to_local(global_pixels, npix_submap, global2local):
+    """
+    Convert global pixel indices to local submaps and pixels within the submap.
+
+    Args:
+        global_pixels (array):  The global pixel indices (size nsamples).
+        npix_submap (int):  The number of pixels in each submap.
+        global2local (array, int64):  The local submap for each global submap.
+
+    Returns:
+        (tuple of array):  The (local submap, pixel within submap) for each global pixel (each of size nsamples).
+    """
+    quotient, remainder = jnp.divmod(global_pixels, npix_submap)
+    local_pixels = jnp.where(global_pixels < 0, -1, remainder)
+    local_submaps = jnp.where(global_pixels < 0, -1, global2local[quotient])
+    return (local_submaps, local_pixels)
+
+
+def scan_map_inner(
+    mapdata,
+    npix_submap,
+    global2local,
+    pixels,
+    weights,
+    det_data,
+    data_scale,
+    should_zero,
+    should_subtract,
+    should_scale,
+):
+    """
+    Applies scan_map to a given interval.
+
+    Args:
+        mapdata (array, ?):  The local piece of the map (size ?*npix_submap*nmap).
+        npix_submap (int):  The number of pixels in each submap.
+        global2local (array, int64):  The local submap for each global submap.
+        pixels (array, int): pixels (size nsample)
+        weights (array, float64):  The pointing matrix weights (size: nsample*nmap).
+        det_data (array, float64):  The timestream on which to accumulate the map values (size nsample).
+        data_scale (float): unit rescaling
+        should_zero (bool): should we zero det_data
+        should_subtract (bool): should we subtract from det_data
+        should_scale (bool): should we scale the detector data by the map values
+
+    Returns:
+        det_data
+    """
+    # Get local submap and pixels
+    submap, subpix = global_to_local(pixels, npix_submap, global2local)
+    # by default JAX will put any value where the indices were invalid instead of erroring out
+    mapdata = mapdata[submap, subpix, :]
+
+    # computes the update term
+    update = jnp.sum(mapdata * weights, axis=1) * data_scale
+
+    # updates det_data
+    if should_zero:
+        det_data = jnp.zeros_like(det_data)
+    if should_subtract:
+        new_det_data = det_data - update
+    elif should_scale:
+        new_det_data = det_data * update
+    else:
+        new_det_data = det_data + update
+
+    # mask invalid values and returns
+    valid_samples = (subpix >= 0) & (submap >= 0)
+    return jnp.where(valid_samples, new_det_data, det_data)
+
+
+# maps over intervals and detectors
+# scan_map_inner = jax_xmap(scan_map_inner,
+#                              in_axes=[[...], # mapdata
+#                                       [...], # npix_submap
+#                                       [...], # global2local
+#                                       ['detectors','intervals',...], # pixels
+#                                       ['detectors','intervals',...], # weights
+#                                       ['detectors','intervals',...], # det_data
+#                                       [...], # data_scale
+#                                       [...], # should_zero
+#                                       [...]], # should_subtract
+#                                       [...]], # should_scale
+#                              out_axes=['detectors','intervals',...])
+# TODO xmap is commented out for now due to a [bug with static argnum](https://github.com/google/jax/issues/10741)
+scan_map_inner = jax.vmap(
+    scan_map_inner,
+    in_axes=[None, None, None, 0, 0, 0, None, None, None, None],
+    out_axes=0,
+)  # loop on intervals
+scan_map_inner = jax.vmap(
+    scan_map_inner,
+    in_axes=[None, None, None, 0, 0, 0, None, None, None, None],
+    out_axes=0,
+)  # loop on detectors
+
+
+def scan_map_interval(
+    mapdata,
+    npix_submap,
+    global2local,
+    det_data,
+    det_data_index,
+    pixels,
+    pixels_index,
+    weights,
+    weight_index,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+    data_scale,
+    should_zero,
+    should_subtract,
+    should_scale,
+):
+    """
+    Process all the intervals as a block.
+
+    Args:
+        mapdata (array, ?):  The local piece of the map (size ?*npix_submap*nmap).
+        npix_submap (int):  The number of pixels in each submap.
+        global2local (array, int64):  The local submap for each global submap.
+        det_data (array, float): size ???*n_samp
+        det_data_index (array, int): The indexes of the det_data (size n_det)
+        pixels (array, int): pixels (size ???*n_samp)
+        pixels_index (array, int): The indexes of the pixels (size n_det)
+        weights (optional array, float64): The flat packed detectors weights for the specified mode (size ???*n_samp*3)
+        weight_index (optional array, int): The indexes of the weights (size n_det)
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+        data_scale (float): unit scaling
+        should_zero (bool): should we zero det_data
+        should_subtract (bool): should we subtract from det_data
+        should_scale (bool): should we scale the detector data by the map values
+
+    Returns:
+        det_data (array, float): size ???*n_samp
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"scan_map: jit-compiling.")
+
+    # extract interval slices
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    pixels_interval = JaxIntervals.get(
+        pixels, (pixels_index, intervals)
+    )  # pixels[pixels_index, intervals]
+    det_data_interval = JaxIntervals.get(
+        det_data, (det_data_index, intervals)
+    )  # det_data[det_data_index, intervals]
+    if weight_index is None:
+        weights_interval = jnp.ones_like(pixels_interval)
+    elif weights.ndim == 2:
+        weights_interval = JaxIntervals.get(
+            weights, (weight_index, intervals)
+        )  # weights[weight_index, intervals]
+    else:
+        weights_interval = JaxIntervals.get(
+            weights, (weight_index, intervals, ALL)
+        )  # weights[weight_index, intervals, :]
+
+    # does the computation
+    new_det_data_interval = scan_map_inner(
+        mapdata,
+        npix_submap,
+        global2local,
+        pixels_interval,
+        weights_interval,
+        det_data_interval,
+        data_scale,
+        should_zero,
+        should_subtract,
+        should_scale,
+    )
+
+    # updates results and returns
+    # det_data[det_data_index, intervals] = new_det_data_interval
+    det_data = JaxIntervals.set(
+        det_data, (det_data_index, intervals), new_det_data_interval
+    )
+    return det_data
+
+
+# jit compiling
+scan_map_interval = jax.jit(
+    scan_map_interval,
+    static_argnames=[
+        "npix_submap",
+        "intervals_max_length",
+        "should_zero",
+        "should_subtract",
+        "should_scale",
+    ],
+    donate_argnums=[3],
+)  # donates det_data
+
+
+@kernel(impl=ImplementationType.JAX, name="scan_map")
+def scan_map_jax(
+    global2local,
+    n_pix_submap,
+    mapdata,
+    det_data,
+    det_data_index,
+    pixels,
+    pixels_index,
+    weights,
+    weight_index,
+    intervals,
+    data_scale=1.0,
+    should_zero=False,
+    should_subtract=False,
+    should_scale=False,
+    use_accel=False,
+):
+    """
+    Kernel for scanning a map into timestreams.
+
+    This uses a local piece of a distributed map and the pointing matrix for local
+    detectors to generate timestream values.
+
+    Args:
+        global2local (array):  The mapping from global submap to local submap index.
+        n_pix_submap (int):  The number of pixels per submap.
+        mapdata (Pixels):  The local piece of the map.
+        det_data (array, float): size ???*n_samp
+        det_data_index (array, int): The indexes of the det_data (size n_det)
+        pixels (array, int): pixels (size ???*n_samp)
+        pixels_index (array, int): The indexes of the pixels (size n_det)
+        weights (optional array, float64): The flat packed detectors weights for the specified mode (size ???*n_samp*3)
+        weight_index (optional array, int): The indexes of the weights (size n_det)
+        intervals (array, Interval): The intervals to modify (size n_view)
+        map_dist (PixelDistribution): encapsulate information to translate the pixel mapping
+        data_scale (float): unit scalling
+        should_zero (bool): should we zero det_data
+        should_subtract (bool): should we subtract from det_data
+        use_accel (bool): should we use an accelerator
+
+    Returns:
+        None: det_data is updated in place
+    """
+    # prepares inputs
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    mapdata = MutableJaxArray.to_array(mapdata.data)
+    global2local = MutableJaxArray.to_array(global2local)
+    det_data_input = MutableJaxArray.to_array(det_data)
+    det_data_index = MutableJaxArray.to_array(det_data_index)
+    pixels = MutableJaxArray.to_array(pixels)
+    pixels_index = MutableJaxArray.to_array(pixels_index)
+    weights = MutableJaxArray.to_array(weights)
+    weight_index = MutableJaxArray.to_array(weight_index)
+
+    # performs computation and updates det_data in place
+    det_data[:] = scan_map_interval(
+        mapdata,
+        n_pix_submap,
+        global2local,
+        det_data_input,
+        det_data_index,
+        pixels,
+        pixels_index,
+        weights,
+        weight_index,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+        data_scale,
+        should_zero,
+        should_subtract,
+        should_scale,
+    )
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_mapmaker_solve", "ops_scan_map")'
diff --git a/src/toast/ops/scan_map/kernels_numpy.py b/src/toast/ops/scan_map/kernels_numpy.py
new file mode 100644
index 000000000..407dc9901
--- /dev/null
+++ b/src/toast/ops/scan_map/kernels_numpy.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="scan_map")
+def scan_map_numpy(
+    global2local,
+    n_pix_submap,
+    mapdata,
+    det_data,
+    det_data_index,
+    pixels,
+    pixels_index,
+    weights,
+    weight_index,
+    intervals,
+    data_scale=1.0,
+    should_zero=False,
+    should_subtract=False,
+    should_scale=False,
+    use_accel=False,
+):
+    nmap = weights.shape[-1]
+    local_map = mapdata.reshape((-1, nmap))
+    for idet in range(len(det_data_index)):
+        pix = pixels[pixels_index[idet]]
+        wts = weights[weight_index[idet]]
+        tod = det_data[det_data_index[idet]]
+        for view in intervals:
+            vslice = slice(view.first, view.last + 1)
+
+            # Get local submaps and pixel indices within each submap
+            good = pix[vslice] >= 0
+            local_submap_pix = pix[vslice] % n_pix_submap
+            local_submap = global2local[pix[vslice] // n_pix_submap]
+
+            # Indices into the local buffer
+            map_pixels = local_submap * n_pix_submap + local_submap_pix
+
+            # Local scanned TOD
+            local_tod = data_scale * local_map[map_pixels[good]] * wts[vslice][good]
+
+            # Accumulate
+            if should_zero:
+                tod[vslice] = 0
+            if should_subtract:
+                tod[vslice][good] -= local_tod
+            elif should_scale:
+                tod[vslice][good] *= local_tod
+            else:
+                tod[vslice][good] += local_tod
diff --git a/src/toast/ops/scan_map.py b/src/toast/ops/scan_map/scan_map.py
similarity index 64%
rename from src/toast/ops/scan_map.py
rename to src/toast/ops/scan_map/scan_map.py
index 1b666d1ee..bdd322863 100644
--- a/src/toast/ops/scan_map.py
+++ b/src/toast/ops/scan_map/scan_map.py
@@ -6,13 +6,14 @@
 import traitlets
 from astropy import units as u
 
-from .._libtoast import scan_map_float32, scan_map_float64
-from ..observation import default_values as defaults
-from ..pixels import PixelData, PixelDistribution
-from ..timing import function_timer
-from ..traits import Bool, Int, Unicode, Unit, trait_docs
-from ..utils import AlignedF64, Logger, unit_conversion
-from .operator import Operator
+from ...accelerator import ImplementationType
+from ...observation import default_values as defaults
+from ...pixels import PixelData, PixelDistribution
+from ...timing import function_timer
+from ...traits import Bool, Int, Unicode, Unit, UseEnum, trait_docs
+from ...utils import AlignedF64, Logger, unit_conversion
+from ..operator import Operator
+from .kernels import ImplementationType, scan_map
 
 
 @trait_docs
@@ -65,9 +66,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         # Check that the detector data is set
         if self.det_data is None:
             raise RuntimeError("You must set the det_data trait before calling exec()")
@@ -106,86 +110,47 @@ def _exec(self, data, detectors=None, **kwargs):
                     raise RuntimeError(msg)
 
             # If our output detector data does not yet exist, create it
-            exists = ob.detdata.ensure(
-                self.det_data, detectors=dets, create_units=self.det_data_units
+            ob.detdata.ensure(
+                self.det_data,
+                detectors=dets,
+                create_units=self.det_data_units,
+                accel=use_accel,
             )
 
+            intervals = ob.intervals[self.view].data
+            det_data = ob.detdata[self.det_data].data
+            det_data_indx = ob.detdata[self.det_data].indices(dets)
+            pixels = ob.detdata[self.pixels].data
+            pixels_indx = ob.detdata[self.pixels].indices(dets)
             data_scale = unit_conversion(
                 map_data.units, ob.detdata[self.det_data].units
             )
-
-            views = ob.view[self.view]
-            for ivw, vw in enumerate(views):
-                view_samples = None
-                if vw.start is None:
-                    # This is a view of the whole obs
-                    view_samples = ob.n_local_samples
-                else:
-                    view_samples = vw.stop - vw.start
-
-                # Temporary array, re-used for all detectors
-                maptod_raw = AlignedF64.zeros(view_samples)
-                maptod = maptod_raw.array()
-
-                for det in dets:
-                    # The pixels, weights, and data.
-                    pix = views.detdata[self.pixels][ivw][det]
-                    if self.weights is None:
-                        wts = np.ones(pix.size, dtype=np.float64)
-                    else:
-                        wts = views.detdata[self.weights][ivw][det]
-                    ddata = views.detdata[self.det_data][ivw][det]
-
-                    # Get local submap and pixels
-                    local_sm, local_pix = map_dist.global_pixel_to_submap(pix)
-
-                    # We support projecting from either float64 or float32 maps.
-
-                    maptod[:] = 0.0
-
-                    if map_data.dtype.char == "d":
-                        scan_map_float64(
-                            map_data.distribution.n_pix_submap,
-                            map_data.n_value,
-                            local_sm.astype(np.int64),
-                            local_pix.astype(np.int64),
-                            map_data.raw,
-                            wts.astype(np.float64).reshape(-1),
-                            maptod,
-                        )
-                    elif map_data.dtype.char == "f":
-                        scan_map_float32(
-                            map_data.distribution.n_pix_submap,
-                            map_data.n_value,
-                            local_sm.astype(np.int64),
-                            local_pix.astype(np.int64),
-                            map_data.raw,
-                            wts.astype(np.float64).reshape(-1),
-                            maptod,
-                        )
-                    else:
-                        raise RuntimeError(
-                            "Projection supports only float32 and float64 binned maps"
-                        )
-
-                    maptod *= data_scale
-
-                    # zero-out if needed
-                    if self.zero:
-                        ddata[:] = 0.0
-
-                    # Add or subtract.  Note that the map scanned timestream will have
-                    # zeros anywhere that the pointing is bad, but those samples (and
-                    # any other detector flags) should be handled at other steps of the
-                    # processing.
-                    if self.subtract:
-                        ddata[:] -= maptod
-                    else:
-                        ddata[:] += maptod
-
-                del maptod
-                maptod_raw.clear()
-                del maptod_raw
+            if self.weights is None:
+                # Use empty arrays, rather than None, so that we can pass that more
+                # easily to compiled kernels that expect a buffer.
+                weights = np.array([], dtype=np.float64)
+                weight_indx = np.array([], dtype=np.int32)
+            else:
+                weights = ob.detdata[self.weights].data
+                weight_indx = ob.detdata[self.weights].indices(dets)
+            scan_map(
+                map_dist.global_submap_to_local,
+                map_dist.n_pix_submap,
+                map_data.data,
+                det_data,
+                det_data_indx,
+                pixels,
+                pixels_indx,
+                weights,
+                weight_indx,
+                intervals,
+                data_scale=data_scale,
+                should_zero=self.zero,
+                should_subtract=self.subtract,
+                should_scale=False,
+                impl=implementation,
+                use_accel=use_accel,
+            )
 
         return
 
@@ -210,6 +175,17 @@ def _provides(self):
         prov = {"meta": list(), "shared": list(), "detdata": list()}
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
+    def _supports_accel(self):
+        return True
+
 
 @trait_docs
 class ScanMask(Operator):
@@ -256,9 +232,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         # Check that the detector data is set
         if self.det_flags is None:
             raise RuntimeError("You must set the det_flags trait before calling exec()")
@@ -308,6 +287,7 @@ def _finalize(self, data, **kwargs):
 
     def _requires(self):
         req = {
+            "meta": list(),
             "global": [self.mask_key],
             "shared": list(),
             "detdata": [self.pixels, self.det_flags],
@@ -318,9 +298,13 @@ def _requires(self):
         return req
 
     def _provides(self):
+        # TODO shouldn't self.det_flags be in provide rather than require since it can be created here?
         prov = {"meta": list(), "shared": list(), "detdata": list()}
         return prov
 
+    def _supports_accel(self):
+        return False
+
 
 @trait_docs
 class ScanScale(Operator):
@@ -360,9 +344,12 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     @function_timer
-    def _exec(self, data, detectors=None, **kwargs):
+    def _exec(self, data, detectors=None, use_accel=False, **kwargs):
         log = Logger.get()
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         # Check that the detector data is set
         if self.det_data is None:
             raise RuntimeError("You must set the det_data trait before calling exec()")
@@ -398,64 +385,37 @@ def _exec(self, data, detectors=None, **kwargs):
                 log.error(msg)
                 raise RuntimeError(msg)
 
-            views = ob.view[self.view]
-            for ivw, vw in enumerate(views):
-                view_samples = None
-                if vw.start is None:
-                    # This is a view of the whole obs
-                    view_samples = ob.n_local_samples
-                else:
-                    view_samples = vw.stop - vw.start
-
-                # Temporary array, re-used for all detectors
-                maptod_raw = AlignedF64.zeros(view_samples)
-                maptod = maptod_raw.array()
-
-                for det in dets:
-                    # The pixels, weights, and data.
-                    pix = views.detdata[self.pixels][ivw][det]
-                    ddata = views.detdata[self.det_data][ivw][det]
-
-                    # Get local submap and pixels
-                    local_sm, local_pix = map_dist.global_pixel_to_submap(pix)
-
-                    # We support projecting from either float64 or float32 maps.  We
-                    # use a shortcut here by passing the original timestream values
-                    # as the pointing "weights", so that the output is equal to the
-                    # pixel values times the original timestream.
-
-                    maptod[:] = 0.0
-
-                    if map_data.dtype.char == "d":
-                        scan_map_float64(
-                            map_data.distribution.n_pix_submap,
-                            1,
-                            local_sm.astype(np.int64),
-                            local_pix.astype(np.int64),
-                            map_data.raw,
-                            ddata.astype(np.float64).reshape(-1),
-                            maptod,
-                        )
-                    elif map_data.dtype.char == "f":
-                        scan_map_float32(
-                            map_data.distribution.n_pix_submap,
-                            1,
-                            local_sm.astype(np.int64),
-                            local_pix.astype(np.int64),
-                            map_data.raw,
-                            ddata.astype(np.float64).reshape(-1),
-                            maptod,
-                        )
-                    else:
-                        raise RuntimeError(
-                            "Projection supports only float32 and float64 binned maps"
-                        )
-
-                    ddata[:] = maptod
-
-                del maptod
-                maptod_raw.clear()
-                del maptod_raw
+            intervals = ob.intervals[self.view].data
+            det_data = ob.detdata[self.det_data].data
+            det_data_indx = ob.detdata[self.det_data].indices(dets)
+            pixels = ob.detdata[self.pixels].data
+            pixels_indx = ob.detdata[self.pixels].indices(dets)
+            if self.weights is None:
+                # Use empty arrays, rather than None, so that we can pass that more
+                # easily to compiled kernels that expect a buffer.
+                weights = np.array([], dtype=np.float64)
+                weight_indx = np.array([], dtype=np.int32)
+            else:
+                weights = ob.detdata[self.weights].data
+                weight_indx = ob.detdata[self.weights].indices(dets)
+            scan_map(
+                map_dist.global_submap_to_local,
+                map_dist.n_pix_submap,
+                map_data.data,
+                det_data,
+                det_data_indx,
+                pixels,
+                pixels_indx,
+                weights,
+                weight_indx,
+                intervals,
+                data_scale=1.0,
+                should_zero=False,
+                should_subtract=False,
+                should_scale=True,
+                impl=implementation,
+                use_accel=use_accel,
+            )
 
         return
 
@@ -464,6 +424,7 @@ def _finalize(self, data, **kwargs):
 
     def _requires(self):
         req = {
+            "meta": list(),
             "global": [self.map_key],
             "shared": list(),
             "detdata": [self.pixels, self.weights, self.det_data],
diff --git a/src/toast/ops/stokes_weights/CMakeLists.txt b/src/toast/ops/stokes_weights/CMakeLists.txt
new file mode 100644
index 000000000..1f476a6e1
--- /dev/null
+++ b/src/toast/ops/stokes_weights/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    stokes_weights.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/ops/stokes_weights
+)
diff --git a/src/toast/ops/stokes_weights/__init__.py b/src/toast/ops/stokes_weights/__init__.py
new file mode 100644
index 000000000..fbd584943
--- /dev/null
+++ b/src/toast/ops/stokes_weights/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .stokes_weights import StokesWeights
diff --git a/src/toast/ops/stokes_weights/kernels.py b/src/toast/ops/stokes_weights/kernels.py
new file mode 100644
index 000000000..dc38ab59d
--- /dev/null
+++ b/src/toast/ops/stokes_weights/kernels.py
@@ -0,0 +1,100 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ... import qarray as qa
+from ..._libtoast import stokes_weights_I as libtoast_stokes_weights_I
+from ..._libtoast import stokes_weights_IQU as libtoast_stokes_weights_IQU
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import stokes_weights_I_numpy, stokes_weights_IQU_numpy
+
+if use_accel_jax:
+    from .kernels_jax import stokes_weights_I_jax, stokes_weights_IQU_jax
+
+
+@kernel(impl=ImplementationType.COMPILED, name="stokes_weights_I")
+def stokes_weights_I_compiled(*args, use_accel=False):
+    return libtoast_stokes_weights_I(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="stokes_weights_IQU")
+def stokes_weights_IQU_compiled(*args, use_accel=False):
+    return libtoast_stokes_weights_IQU(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def stokes_weights_I(
+    weight_index,
+    weights,
+    intervals,
+    cal,
+    use_accel=False,
+):
+    """Kernel for computing trivial intensity-only Stokes pointing weights.
+
+    Args:
+        weight_index (array):  The index into the weights array for each detector.
+        weights (array):  The array of I, Q, and U weights at each sample for each
+            detector.
+        intervals (array):  The array of sample intervals.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return stokes_weights_I(
+        weight_index,
+        weights,
+        intervals,
+        cal,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def stokes_weights_IQU(
+    quat_index,
+    quats,
+    weight_index,
+    weights,
+    hwp,
+    intervals,
+    epsilon,
+    cal,
+    use_accel=False,
+):
+    """Kernel for computing the I/Q/U Stokes pointing weights.
+
+    Args:
+        quat_index (array):  The index into the detector quaternion array for each
+            detector.
+        quats (array):  The array of detector quaternions for each sample.
+        weight_index (array):  The index into the weights array for each detector.
+        weights (array):  The array of I, Q, and U weights at each sample for each
+            detector.
+        hwp (array):  The array of orientation angles for an ideal half wave plate.
+        intervals (array):  The array of sample intervals.
+        epsilon (float):  The cross polar leakage.
+        cal (float):  Calibration factor to apply to the weights.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return stokes_weights_IQU(
+        quat_index,
+        quats,
+        weight_index,
+        weights,
+        hwp,
+        intervals,
+        epsilon,
+        cal,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
diff --git a/src/toast/ops/stokes_weights/kernels_jax.py b/src/toast/ops/stokes_weights/kernels_jax.py
new file mode 100644
index 000000000..84227b8c8
--- /dev/null
+++ b/src/toast/ops/stokes_weights/kernels_jax.py
@@ -0,0 +1,218 @@
+# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+from jax.experimental.maps import xmap as jax_xmap
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import ALL, INTERVALS_JAX, JaxIntervals
+from ...jax.math import qarray
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def stokes_weights_IQU_inner(eps, cal, pin, hwpang):
+    """
+    Compute the Stokes weights for one detector.
+
+    Args:
+        eps (float):  The cross polar response.
+        cal (float):  A constant to apply to the pointing weights.
+        pin (array, float64):  The array of detector quaternions (size 4).
+        hwpang (float64):  The HWP angle (could be None).
+
+    Returns:
+        weights (array, float64):  The detector weights for the specified mode (size 3)
+    """
+    # applies quaternion rotations
+    zaxis = jnp.array([0.0, 0.0, 1.0])
+    dir = qarray.rotate_one_one(pin, zaxis)
+    xaxis = jnp.array([1.0, 0.0, 0.0])
+    orient = qarray.rotate_one_one(pin, xaxis)
+
+    # computes by and bx
+    by = orient[0] * dir[1] - orient[1] * dir[0]
+    bx = (
+        orient[0] * (-dir[2] * dir[0])
+        + orient[1] * (-dir[2] * dir[1])
+        + orient[2] * (dir[0] * dir[0] + dir[1] * dir[1])
+    )
+
+    # computes detang
+    detang = jnp.arctan2(by, bx)
+    detang = detang + 2.0 * hwpang
+    detang = 2.0 * detang
+
+    # puts values into weights
+    eta = (1.0 - eps) / (1.0 + eps)
+    weights = jnp.array([cal, jnp.cos(detang) * eta * cal, jnp.sin(detang) * eta * cal])
+    return weights
+
+
+# maps over samples, intervals and detectors
+# stokes_weights_IQU_inner = jax_xmap(
+#    stokes_weights_IQU_inner,
+#    in_axes=[
+#        ["detectors"],  # epsilon
+#        [...],  # cal
+#        ["detectors", "intervals", "interval_size", ...],  # quats
+#        ["intervals", "interval_size"],
+#    ],  # hwp
+#    out_axes=["detectors", "intervals", "interval_size", ...],
+# )
+# using vmap as the static arguments triggers the following error:
+# "ShardingContext cannot be used with xmap"
+# TODO revisit once this issue is solved [bug with static argnum](https://github.com/google/jax/issues/10741)
+stokes_weights_IQU_inner = jax.vmap(
+    stokes_weights_IQU_inner, in_axes=[None, None, 0, 0], out_axes=0
+)  # interval_size
+stokes_weights_IQU_inner = jax.vmap(
+    stokes_weights_IQU_inner, in_axes=[None, None, 0, 0], out_axes=0
+)  # intervals
+stokes_weights_IQU_inner = jax.vmap(
+    stokes_weights_IQU_inner, in_axes=[0, None, 0, None], out_axes=0
+)  # detectors
+
+
+def stokes_weights_IQU_interval(
+    quat_index,
+    quats,
+    weight_index,
+    weights,
+    hwp,
+    epsilon,
+    cal,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+):
+    """
+    Process all the intervals as a block.
+
+    Args:
+        quat_index (array, int): size n_det
+        quats (array, double): size ???*n_samp*4
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, float64): The flat packed detectors weights for the specified mode (size n_det*n_samp*3)
+        hwp (array, float64):  The HWP angles (size n_samp).
+        epsilon (array, float):  The cross polar response (size n_det).
+        cal (float):  A constant to apply to the pointing weights.
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+
+    Returns:
+        weights
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"stokes_weights_IQU: jit-compiling.")
+
+    # extract interval slices
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    quats_interval = JaxIntervals.get(
+        quats, (quat_index, intervals, ALL)
+    )  # quats[quat_index,intervals,:]
+
+    # insures hwp is a non empty array
+    if hwp.size == 0:
+        nb_intervals = interval_starts.size
+        hwp_interval = jnp.zeros((nb_intervals, intervals_max_length))
+    else:
+        hwp_interval = JaxIntervals.get(hwp, intervals)  # hwp[intervals]
+
+    # does the computation
+    new_weights_interval = stokes_weights_IQU_inner(
+        epsilon, cal, quats_interval, hwp_interval
+    )
+
+    # updates results and returns
+    # weights[weight_index,intervals,:] = new_weights_interval
+    weights = JaxIntervals.set(
+        weights, (weight_index, intervals, ALL), new_weights_interval
+    )
+    return weights
+
+
+# jit compiling
+stokes_weights_IQU_interval = jax.jit(
+    stokes_weights_IQU_interval,
+    static_argnames=["cal", "intervals_max_length"],
+    donate_argnums=[3],
+)  # donates weights
+
+
+@kernel(impl=ImplementationType.JAX, name="stokes_weights_IQU")
+def stokes_weights_IQU_jax(
+    quat_index, quats, weight_index, weights, hwp, intervals, epsilon, cal, use_accel
+):
+    """
+    Compute the Stokes weights for the "IQU" mode.
+
+    Args:
+        quat_index (array, int): size n_det
+        quats (array, double): size ???*n_samp*4
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, float64): The flat packed detectors weights for the specified mode (size n_det*n_samp*3)
+        hwp (array, float64):  The HWP angles (size n_samp).
+        intervals (array, Interval): The intervals to modify (size n_view)
+        epsilon (array, float):  The cross polar response (size n_det).
+        cal (float):  A constant to apply to the pointing weights.
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in weights).
+    """
+    # prepares inputs
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    quat_index_input = MutableJaxArray.to_array(quat_index)
+    quats_input = MutableJaxArray.to_array(quats)
+    weight_index_input = MutableJaxArray.to_array(weight_index)
+    weights_input = MutableJaxArray.to_array(weights)
+    hwp_input = MutableJaxArray.to_array(hwp)
+    epsilon_input = MutableJaxArray.to_array(epsilon)
+
+    # runs computation
+    weights[:] = stokes_weights_IQU_interval(
+        quat_index_input,
+        quats_input,
+        weight_index_input,
+        weights_input,
+        hwp_input,
+        epsilon_input,
+        cal,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+    )
+
+
+@kernel(impl=ImplementationType.JAX, name="stokes_weights_I")
+def stokes_weights_I_jax(weight_index, weights, intervals, cal, use_accel):
+    """
+    Compute the Stokes weights for the "I" mode.
+    NOTE this does not use JAX as there is too little computation
+
+    Args:
+        weight_index (array, int): The indexes of the weights (size n_det)
+        weights (array, float64): The flat packed detectors weights for the specified mode (size n_det*n_samp)
+        intervals (array, Interval): The intervals to modify (size n_view)
+        cal (float):  A constant to apply to the pointing weights.
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in weights).
+    """
+    # iterate on the intervals
+    for interval in intervals:
+        interval_start = interval.first
+        interval_end = interval.last + 1
+        weights[weight_index, interval_start:interval_end] = cal
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("ops_pointing_healpix"); toast.tests.run("ops_sim_tod_dipole")'
diff --git a/src/toast/ops/stokes_weights/kernels_numpy.py b/src/toast/ops/stokes_weights/kernels_numpy.py
new file mode 100644
index 000000000..fdcaa0a0f
--- /dev/null
+++ b/src/toast/ops/stokes_weights/kernels_numpy.py
@@ -0,0 +1,69 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ... import qarray as qa
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="stokes_weights_IQU")
+def stokes_weights_IQU_numpy(
+    quat_index, quats, weight_index, weights, hwp, intervals, epsilon, cal, use_accel
+):
+    if hwp is not None and len(hwp) == 0:
+        hwp = None
+    zaxis = np.array([0, 0, 1], dtype=np.float64)
+    xaxis = np.array([1, 0, 0], dtype=np.float64)
+    for idet in range(len(quat_index)):
+        qidx = quat_index[idet]
+        widx = weight_index[idet]
+        eta = (1.0 - epsilon[idet]) / (1.0 + epsilon[idet])
+        for vw in intervals:
+            samples = slice(vw.first, vw.last + 1, 1)
+            dir = qa.rotate(quats[qidx][samples], zaxis)
+            orient = qa.rotate(quats[qidx][samples], xaxis)
+
+            # The vector orthogonal to the line of sight that is parallel
+            # to the local meridian.
+            dir_ang = np.arctan2(dir[:, 1], dir[:, 0])
+            dir_r = np.sqrt(1.0 - dir[:, 2] * dir[:, 2])
+            m_z = dir_r
+            m_x = -dir[:, 2] * np.cos(dir_ang)
+            m_y = -dir[:, 2] * np.sin(dir_ang)
+
+            # Compute the rotation angle from the meridian vector to the
+            # orientation vector.  The direction vector is normal to the plane
+            # containing these two vectors, so the rotation angle is:
+            #
+            # angle = atan2((v_m x v_o) . v_d, v_m . v_o)
+            # angle = atan2(
+            #     d_x (m_y o_z - m_z o_y)
+            #       - d_y (m_x o_z - m_z o_x)
+            #       + d_z (m_x o_y - m_y o_x),
+            #     m_x o_x + m_y o_y + m_z o_z
+            # )
+            #
+            ay = (
+                dir[:, 0] * (m_y * orient[:, 2] - m_z * orient[:, 1])
+                - dir[:, 1] * (m_x * orient[:, 2] - m_z * orient[:, 0])
+                + dir[:, 2] * (m_x * orient[:, 1] - m_y * orient[:, 0])
+            )
+            ax = m_x * orient[:, 0] + m_y * orient[:, 1] + m_z * orient[:, 2]
+            ang = np.arctan2(ay, ax)
+            if hwp is not None:
+                ang += 2.0 * hwp[samples]
+            ang *= 2.0
+            weights[widx][samples, 0] = cal
+            weights[widx][samples, 1] = cal * eta * np.cos(ang)
+            weights[widx][samples, 2] = cal * eta * np.sin(ang)
+
+
+@kernel(impl=ImplementationType.NUMPY, name="stokes_weights_I")
+def stokes_weights_I_numpy(weight_index, weights, intervals, cal, use_accel):
+    for idet in range(len(weight_index)):
+        widx = weight_index[idet]
+        for vw in intervals:
+            samples = slice(vw.first, vw.last + 1, 1)
+            weights[widx][samples] = cal
diff --git a/src/toast/ops/stokes_weights.py b/src/toast/ops/stokes_weights/stokes_weights.py
similarity index 67%
rename from src/toast/ops/stokes_weights.py
rename to src/toast/ops/stokes_weights/stokes_weights.py
index 7f7535ceb..549c61d32 100644
--- a/src/toast/ops/stokes_weights.py
+++ b/src/toast/ops/stokes_weights/stokes_weights.py
@@ -5,16 +5,13 @@
 import numpy as np
 import traitlets
 
-from .. import qarray as qa
-from .._libtoast import stokes_weights, stokes_weights_I, stokes_weights_IQU
-from ..healpix import HealpixPixels
-from ..observation import default_values as defaults
-from ..pixels import PixelDistribution
-from ..timing import function_timer
-from ..traits import Bool, Instance, Int, Unicode, trait_docs
-from ..utils import Environment, Logger
-from .delete import Delete
-from .operator import Operator
+from ...accelerator import ImplementationType
+from ...observation import default_values as defaults
+from ...timing import function_timer
+from ...traits import Bool, Instance, Int, Unicode, trait_docs
+from ...utils import Environment, Logger
+from ..operator import Operator
+from .kernels import stokes_weights_I, stokes_weights_IQU
 
 
 @trait_docs
@@ -106,8 +103,6 @@ class StokesWeights(Operator):
 
     IAU = Bool(False, help="If True, use the IAU convention rather than COSMO")
 
-    use_python = Bool(False, help="If True, use python implementation")
-
     @traitlets.validate("detector_pointing")
     def _check_detector_pointing(self, proposal):
         detpointing = proposal["value"]
@@ -148,12 +143,12 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
 
         self._nnz = len(self.mode)
 
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         if self.detector_pointing is None:
             raise RuntimeError("The detector_pointing trait must be set")
 
-        if self.use_python and use_accel:
-            raise RuntimeError("Cannot use accelerator with pure python implementation")
-
         # Expand detector pointing
         if self.quats is not None:
             quats_name = self.quats
@@ -246,45 +241,32 @@ def _exec(self, data, detectors=None, use_accel=False, **kwargs):
                 for idet, d in enumerate(dets):
                     det_epsilon[idet] = focalplane[d]["pol_leakage"]
 
-            if self.use_python:
-                hwp_data = None
-                if self.hwp_angle is not None:
+            if self.mode == "IQU":
+                if self.hwp_angle is None:
+                    hwp_data = np.zeros((0,), dtype=np.float64)
+                else:
                     hwp_data = ob.shared[self.hwp_angle].data
-                self._py_stokes_weights(
+                stokes_weights_IQU(
                     quat_indx,
                     ob.detdata[quats_name].data,
                     weight_indx,
                     ob.detdata[self.weights].data,
+                    hwp_data,
                     ob.intervals[self.view].data,
-                    cal,
                     det_epsilon,
-                    hwp_data,
+                    cal,
+                    impl=implementation,
+                    use_accel=use_accel,
                 )
             else:
-                if self.mode == "IQU":
-                    if self.hwp_angle is None:
-                        hwp_data = np.zeros((0,), dtype=np.float64)
-                    else:
-                        hwp_data = ob.shared[self.hwp_angle].data
-                    stokes_weights_IQU(
-                        quat_indx,
-                        ob.detdata[quats_name].data,
-                        weight_indx,
-                        ob.detdata[self.weights].data,
-                        hwp_data,
-                        ob.intervals[self.view].data,
-                        det_epsilon,
-                        cal,
-                        use_accel,
-                    )
-                else:
-                    stokes_weights_I(
-                        weight_indx,
-                        ob.detdata[self.weights].data,
-                        ob.intervals[self.view].data,
-                        cal,
-                        use_accel,
-                    )
+                stokes_weights_I(
+                    weight_indx,
+                    ob.detdata[self.weights].data,
+                    ob.intervals[self.view].data,
+                    cal,
+                    impl=implementation,
+                    use_accel=use_accel,
+                )
         return
 
     def _finalize(self, data, **kwargs):
@@ -304,69 +286,13 @@ def _provides(self):
         prov = {"detdata": [self.weights]}
         return prov
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
     def _supports_accel(self):
-        return self.detector_pointing.supports_accel()
-
-    def _py_stokes_weights(
-        self,
-        quat_indx,
-        quat_data,
-        weight_indx,
-        weight_data,
-        intr_data,
-        cal,
-        det_epsilon,
-        hwp_data,
-    ):
-        """Internal python implementation for comparison tests."""
-        zaxis = np.array([0, 0, 1], dtype=np.float64)
-        xaxis = np.array([1, 0, 0], dtype=np.float64)
-        if self.mode == "IQU":
-            for idet in range(len(quat_indx)):
-                qidx = quat_indx[idet]
-                widx = weight_indx[idet]
-                eta = (1.0 - det_epsilon[idet]) / (1.0 + det_epsilon[idet])
-                for vw in intr_data:
-                    samples = slice(vw.first, vw.last + 1, 1)
-                    dir = qa.rotate(quat_data[qidx][samples], zaxis)
-                    orient = qa.rotate(quat_data[qidx][samples], xaxis)
-
-                    # The vector orthogonal to the line of sight that is parallel
-                    # to the local meridian.
-                    dir_ang = np.arctan2(dir[:, 1], dir[:, 0])
-                    dir_r = np.sqrt(1.0 - dir[:, 2] * dir[:, 2])
-                    m_z = dir_r
-                    m_x = -dir[:, 2] * np.cos(dir_ang)
-                    m_y = -dir[:, 2] * np.sin(dir_ang)
-
-                    # Compute the rotation angle from the meridian vector to the
-                    # orientation vector.  The direction vector is normal to the plane
-                    # containing these two vectors, so the rotation angle is:
-                    #
-                    # angle = atan2((v_m x v_o) . v_d, v_m . v_o)
-                    # angle = atan2(
-                    #     d_x (m_y o_z - m_z o_y)
-                    #       - d_y (m_x o_z - m_z o_x)
-                    #       + d_z (m_x o_y - m_y o_x),
-                    #     m_x o_x + m_y o_y + m_z o_z
-                    # )
-                    #
-                    ay = (
-                        dir[:, 0] * (m_y * orient[:, 2] - m_z * orient[:, 1])
-                        - dir[:, 1] * (m_x * orient[:, 2] - m_z * orient[:, 0])
-                        + dir[:, 2] * (m_x * orient[:, 1] - m_y * orient[:, 0])
-                    )
-                    ax = m_x * orient[:, 0] + m_y * orient[:, 1] + m_z * orient[:, 2]
-                    ang = np.arctan2(ay, ax)
-                    if hwp_data is not None:
-                        ang += 2.0 * hwp_data[samples]
-                    ang *= 2.0
-                    weight_data[widx][samples, 0] = cal
-                    weight_data[widx][samples, 1] = cal * eta * np.cos(ang)
-                    weight_data[widx][samples, 2] = cal * eta * np.sin(ang)
-        else:
-            for idet in range(len(quat_indx)):
-                widx = weight_indx[idet]
-                for vw in intr_data:
-                    samples = slice(vw.first, vw.last + 1, 1)
-                    weight_data[widx][samples] = cal
+        return True
diff --git a/src/toast/pixels.py b/src/toast/pixels.py
index 1a35a295b..a9a4aaa8a 100644
--- a/src/toast/pixels.py
+++ b/src/toast/pixels.py
@@ -37,12 +37,8 @@
     Logger,
 )
 
-if use_accel_jax:
-    import jax
-    import jax.numpy as jnp
 
-
-class PixelDistribution(object):
+class PixelDistribution(AcceleratorObject):
     """Class representing the distribution of submaps.
 
     This object is used to describe the properties of a pixelization scheme and which
@@ -90,6 +86,8 @@ def __init__(self, n_pix=None, n_submap=1000, local_submaps=None, comm=None):
         self._alltoallv_info = None
         self._all_hit_submaps = None
 
+        super().__init__()
+
     def __eq__(self, other):
         local_eq = True
         if self._n_pix != other._n_pix:
@@ -405,6 +403,39 @@ def alltoallv_info(self):
 
         return self._alltoallv_info
 
+    def _accel_exists(self):
+        if use_accel_omp or use_accel_jax:
+            return accel_data_present(self._glob2loc)
+        else:
+            return False
+
+    def _accel_create(self):
+        if use_accel_omp:
+            accel_data_create(self._glob2loc)
+        elif use_accel_jax:
+            self._glob2loc = accel_data_create(self._glob2loc)
+
+    def _accel_update_device(self):
+        if use_accel_omp:
+            _ = accel_data_update_device(self._glob2loc)
+        elif use_accel_jax:
+            self._glob2loc = accel_data_update_device(self._glob2loc)
+
+    def _accel_update_host(self):
+        if use_accel_omp:
+            _ = accel_data_update_host(self._glob2loc)
+        elif use_accel_jax:
+            self._glob2loc = accel_data_update_host(self._glob2loc)
+
+    def _accel_delete(self):
+        if use_accel_omp:
+            accel_data_delete(self._glob2loc)
+        elif use_accel_jax and self._accel_exists():
+            # insures glob2loc has been properly reset
+            # if we observe that its types is still a GPU types
+            # does NOT move data back from GPU
+            self._glob2loc = self._glob2loc.host_data
+
 
 class PixelData(AcceleratorObject):
     """Distributed map-domain data.
@@ -486,7 +517,6 @@ def __init__(self, dist, dtype, n_value=1, units=u.dimensionless_unscaled):
 
         self.raw = self.storage_class.zeros(self._flatshape)
         self.data = self.raw.array().reshape(self._shape)
-        self.data_jax = None
 
         # Allreduce quantities
         self._all_comm_submap = None
@@ -517,7 +547,8 @@ def clear(self):
 
         """
         if hasattr(self, "data"):
-            del self.data
+            # we keep the attribute to avoid errors in _accel_exists
+            self.data = None
         if hasattr(self, "raw"):
             if self.accel_exists():
                 self.accel_delete()
@@ -1133,7 +1164,7 @@ def _accel_exists(self):
         if use_accel_omp:
             return accel_data_present(self.raw)
         elif use_accel_jax:
-            return accel_data_present(self.data_jax)
+            return accel_data_present(self.data)
         else:
             return False
 
@@ -1141,24 +1172,25 @@ def _accel_create(self):
         if use_accel_omp:
             accel_data_create(self.raw)
         elif use_accel_jax:
-            accel_data_create(self.data_jax)
+            self.data = accel_data_create(self.data)
 
     def _accel_update_device(self):
         if use_accel_omp:
             _ = accel_data_update_device(self.raw)
         elif use_accel_jax:
-            self.data_jax = accel_data_update_device(self.data)
+            self.data = accel_data_update_device(self.data)
 
     def _accel_update_host(self):
         if use_accel_omp:
             _ = accel_data_update_host(self.raw)
         elif use_accel_jax:
-            self.data[:] = accel_data_update_host(self.data_jax)
-            self.data_jax = None
+            self.data = accel_data_update_host(self.data)
 
     def _accel_delete(self):
         if use_accel_omp:
             accel_data_delete(self.raw)
-        elif use_accel_jax:
-            del self.data_jax
-            self.data_jax = None
+        elif use_accel_jax and self._accel_exists():
+            # insures data has been properly reset
+            # if we observe that its types is still a GPU types
+            # does NOT move data back from GPU
+            self.data = self.data.host_data
diff --git a/src/toast/scripts/toast_healpix_coadd.py b/src/toast/scripts/toast_healpix_coadd.py
index 51fdaf080..3c939c2d3 100644
--- a/src/toast/scripts/toast_healpix_coadd.py
+++ b/src/toast/scripts/toast_healpix_coadd.py
@@ -19,9 +19,9 @@
 
 import toast
 from toast import PixelData, PixelDistribution
-from toast._libtoast import cov_apply_diag, cov_eigendecompose_diag
 from toast.covariance import covariance_apply, covariance_invert
 from toast.mpi import MPI, Comm, get_world
+from toast.ops.kernels import cov_apply_diag, cov_eigendecompose_diag
 from toast.pixels_io_healpix import (
     filename_is_fits,
     filename_is_hdf5,
diff --git a/src/toast/scripts/toast_mini.py b/src/toast/scripts/toast_mini.py
index 53aa456f2..1e1cfb148 100644
--- a/src/toast/scripts/toast_mini.py
+++ b/src/toast/scripts/toast_mini.py
@@ -20,6 +20,7 @@
 
 import toast
 import toast.ops
+from toast.accelerator.data_localization import display_datamovement
 from toast.schedule_sim_satellite import create_satellite_schedule
 from toast.scripts.benchmarking_utilities import (
     default_sim_atmosphere,
@@ -109,6 +110,8 @@ def parse_arguments():
     args.max_detector = (
         2054  # Hex-packed 1027 pixels (18 rings) times two dets per pixel.
     )
+    # For debugging:
+    # args.max_detector = 4
     args.width = 10
     args.obs_minutes = 60
     args.num_obs = 4380
@@ -256,7 +259,6 @@ def main():
     log.info_rank("Finished map-making in", comm=world_comm, timer=timer)
 
     # Dump all the timing information
-
     timer.stop()
     timer.clear()
     timer.start()
@@ -269,6 +271,9 @@ def main():
     else:
         timer.stop()
 
+    # display information on GPU data movement
+    display_datamovement()
+
 
 if __name__ == "__main__":
     world, procs, rank = toast.mpi.get_world()
diff --git a/src/toast/templates/CMakeLists.txt b/src/toast/templates/CMakeLists.txt
index f2fee4d40..9d6dc859f 100644
--- a/src/toast/templates/CMakeLists.txt
+++ b/src/toast/templates/CMakeLists.txt
@@ -5,9 +5,10 @@ install(FILES
     __init__.py
     template.py
     amplitudes.py
-    offset.py
     fourier2d.py
     subharmonic.py
     gaintemplate.py
     DESTINATION ${PYTHON_SITE}/toast/templates
 )
+
+add_subdirectory(offset)
diff --git a/src/toast/templates/amplitudes.py b/src/toast/templates/amplitudes.py
index fa323d4d7..2c5a09b97 100644
--- a/src/toast/templates/amplitudes.py
+++ b/src/toast/templates/amplitudes.py
@@ -156,7 +156,6 @@ def __init__(
                 self._global_last = self._local_indices[-1]
         self._raw = self._storage_class.zeros(self._n_local)
         self.local = self._raw.array()
-        self.local_jax = None
 
         # Support flagging of template amplitudes.  This can be used to flag some
         # amplitudes if too many timestream samples contributing to the amplitude value
@@ -165,7 +164,7 @@ def __init__(
         # a bit of memory and use a whole byte per amplitude.
         self._raw_flags = AlignedU8.zeros(self._n_local)
         self.local_flags = self._raw_flags.array()
-        self.local_flags_jax = None
+
         super().__init__()
 
     def clear(self):
@@ -177,14 +176,14 @@ def clear(self):
 
         """
         if hasattr(self, "local"):
-            del self.local
+            self.local = None
         if hasattr(self, "_raw"):
             if self.accel_exists():
                 self.accel_delete()
             self._raw.clear()
             del self._raw
         if hasattr(self, "local_flags"):
-            del self.local_flags
+            self.local_flags = None
         if hasattr(self, "_raw_flags"):
             self._raw_flags.clear()
             del self._raw_flags
@@ -665,8 +664,8 @@ def _accel_exists(self):
         if use_accel_omp:
             return accel_data_present(self._raw) and accel_data_present(self._raw_flags)
         elif use_accel_jax:
-            return accel_data_present(self.local_jax) and accel_data_present(
-                self.local_flags_jax
+            return accel_data_present(self.local) and accel_data_present(
+                self.local_flags
             )
         else:
             return False
@@ -676,36 +675,36 @@ def _accel_create(self):
             accel_data_create(self._raw)
             accel_data_create(self._raw_flags)
         elif use_accel_jax:
-            accel_data_create(self.local_jax)
-            accel_data_create(self.local_flags_jax)
+            self.local = accel_data_create(self.local)
+            self.local_flags = accel_data_create(self.local_flags)
 
     def _accel_update_device(self):
         if use_accel_omp:
             _ = accel_data_update_device(self._raw)
             _ = accel_data_update_device(self._raw_flags)
         elif use_accel_jax:
-            self.local_jax = accel_data_update_device(self.local)
-            self.local_flags_jax = accel_data_update_device(self.local_flags)
+            self.local = accel_data_update_device(self.local)
+            self.local_flags = accel_data_update_device(self.local_flags)
 
     def _accel_update_host(self):
         if use_accel_omp:
             _ = accel_data_update_host(self._raw)
             _ = accel_data_update_host(self._raw_flags)
         elif use_accel_jax:
-            self.local[:] = accel_data_update_host(self.local_jax)
-            self.local_flags[:] = accel_data_update_host(self.local_flags_jax)
-            self.local_jax = None
-            self.local_flags_jax = None
+            self.local = accel_data_update_host(self.local)
+            self.local_flags = accel_data_update_host(self.local_flags)
 
     def _accel_delete(self):
         if use_accel_omp:
             accel_data_delete(self._raw)
             accel_data_delete(self._raw_flags)
-        elif use_accel_jax:
-            del self.local_jax
-            del self.local_flags_jax
-            self.local_jax = None
-            self.local_flags_jax = None
+        elif use_accel_jax and self._accel_exists():
+            # insures local and local_flags have been properly reset
+            # if we observe that their types are still GPU types
+            # does NOT move data back from GPU
+            # using self._raw.array() and self._raw_flags.array() should be equivalent
+            self.local = self.local.host_data
+            self.local_flags = self.local_flags.host_data
 
 
 class AmplitudesMap(MutableMapping, AcceleratorObject):
@@ -718,6 +717,7 @@ class AmplitudesMap(MutableMapping, AcceleratorObject):
 
     def __init__(self):
         self._internal = dict()
+        super().__init__()
 
     # Mapping methods
 
diff --git a/src/toast/templates/fourier2d.py b/src/toast/templates/fourier2d.py
index d7a867deb..b2517b5bf 100644
--- a/src/toast/templates/fourier2d.py
+++ b/src/toast/templates/fourier2d.py
@@ -326,7 +326,7 @@ def _zeros(self):
         # z.local_flags[:] = np.where(self._amp_flags, 1, 0)
         return z
 
-    def _add_to_signal(self, detector, amplitudes):
+    def _add_to_signal(self, detector, amplitudes, **kwargs):
         for iob, ob in enumerate(self.data.obs):
             if detector not in ob.local_detectors:
                 continue
@@ -344,7 +344,7 @@ def _add_to_signal(self, detector, amplitudes):
                     1,
                 )
 
-    def _project_signal(self, detector, amplitudes):
+    def _project_signal(self, detector, amplitudes, **kwargs):
         for iob, ob in enumerate(self.data.obs):
             if detector not in ob.local_detectors:
                 continue
@@ -362,7 +362,7 @@ def _project_signal(self, detector, amplitudes):
                     self._templates[iob][ivw][detector],
                 )
 
-    def _add_prior(self, amplitudes_in, amplitudes_out):
+    def _add_prior(self, amplitudes_in, amplitudes_out, **kwargs):
         for iob, ob in enumerate(self.data.obs):
             views = ob.view[self.view]
             for ivw, vw in enumerate(views):
@@ -382,6 +382,6 @@ def _add_prior(self, amplitudes_in, amplitudes_out):
                         mode="same",
                     )
 
-    def _apply_precond(self, amplitudes_in, amplitudes_out):
+    def _apply_precond(self, amplitudes_in, amplitudes_out, **kwargs):
         amplitudes_out.local[:] = amplitudes_in.local
         amplitudes_out.local *= self._norms
diff --git a/src/toast/templates/gaintemplate.py b/src/toast/templates/gaintemplate.py
index 16f71053f..e198294ba 100644
--- a/src/toast/templates/gaintemplate.py
+++ b/src/toast/templates/gaintemplate.py
@@ -7,10 +7,8 @@
 import numpy as np
 import scipy
 
-from ..data import Data
 from ..mpi import MPI
-from ..traits import Bool, Float, Instance, Int, Unicode, trait_docs
-from ..utils import Logger
+from ..traits import Int, Unicode, trait_docs
 from .amplitudes import Amplitudes
 from .template import Template
 
@@ -163,7 +161,7 @@ def _zeros(self):
         # z.local_flags[:] = np.where(self._amp_flags, 1, 0)
         return z
 
-    def _add_to_signal(self, detector, amplitudes):
+    def _add_to_signal(self, detector, amplitudes, **kwargs):
         norder = self.order + 1
         offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
@@ -177,7 +175,7 @@ def _add_to_signal(self, detector, amplitudes):
                 gain_fluctuation = signal_estimate * delta_gain
                 vw[detector] += gain_fluctuation
 
-    def _project_signal(self, detector, amplitudes):
+    def _project_signal(self, detector, amplitudes, **kwargs):
         norder = self.order + 1
         offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
@@ -197,11 +195,11 @@ def _project_signal(self, detector, amplitudes):
                 poly_amps = amplitudes.local[offset : offset + norder]
                 poly_amps += np.dot(LT, vw[detector] * mask)
 
-    def _add_prior(self, amplitudes_in, amplitudes_out):
+    def _add_prior(self, amplitudes_in, amplitudes_out, **kwargs):
         # No prior for this template, nothing to accumulate to output.
         return
 
-    def _apply_precond(self, amplitudes_in, amplitudes_out):
+    def _apply_precond(self, amplitudes_in, amplitudes_out, **kwargs):
         norder = self.order + 1
         for det in self._all_dets:
             offset = self._det_start[det]
diff --git a/src/toast/templates/offset/CMakeLists.txt b/src/toast/templates/offset/CMakeLists.txt
new file mode 100644
index 000000000..942ffe0c0
--- /dev/null
+++ b/src/toast/templates/offset/CMakeLists.txt
@@ -0,0 +1,11 @@
+
+# Install the python files
+
+install(FILES
+    __init__.py
+    offset.py
+    kernels.py
+    kernels_numpy.py
+    kernels_jax.py
+    DESTINATION ${PYTHON_SITE}/toast/templates/offset
+)
diff --git a/src/toast/templates/offset/__init__.py b/src/toast/templates/offset/__init__.py
new file mode 100644
index 000000000..54ce65b49
--- /dev/null
+++ b/src/toast/templates/offset/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+# Namespace imports
+
+from .offset import Offset
diff --git a/src/toast/templates/offset/kernels.py b/src/toast/templates/offset/kernels.py
new file mode 100644
index 000000000..83de48795
--- /dev/null
+++ b/src/toast/templates/offset/kernels.py
@@ -0,0 +1,163 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ..._libtoast import template_offset_add_to_signal as libtoast_offset_add_to_signal
+from ..._libtoast import (
+    template_offset_apply_diag_precond as libtoast_offset_apply_diag_precond,
+)
+from ..._libtoast import (
+    template_offset_project_signal as libtoast_offset_project_signal,
+)
+from ...accelerator import ImplementationType, kernel, use_accel_jax
+from .kernels_numpy import (
+    offset_add_to_signal_numpy,
+    offset_apply_diag_precond_numpy,
+    offset_project_signal_numpy,
+)
+
+if use_accel_jax:
+    from .kernels_jax import (
+        offset_add_to_signal_jax,
+        offset_apply_diag_precond_jax,
+        offset_project_signal_jax,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def offset_add_to_signal(
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    data_index,
+    det_data,
+    intervals,
+    use_accel=False,
+):
+    """Kernel to accumulate offset amplitudes to timestream data.
+
+    Each amplitude value is accumulated to `step_length` number of samples.
+
+    Args:
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int64):  The first amplitude for this detector.
+        n_amp_views (array, int):  The number of amplitudes for each interval.
+        amplitudes (array, double): The amplitude data.
+        data_index (int):  The detector to process.
+        det_data (array, double):  The array of data for all detectors.
+        intervals (array, Interval):  The intervals to process.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return offset_add_to_signal(
+        step_length,
+        amp_offset,
+        n_amp_views,
+        amplitudes,
+        data_index,
+        det_data,
+        intervals,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def offset_project_signal(
+    data_index,
+    det_data,
+    flag_index,
+    flag_data,
+    flag_mask,
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    intervals,
+    use_accel=False,
+):
+    """Kernel to accumulate timestream data into offset amplitudes.
+
+    Chunks of `step_length` number of samples from one detector are accumulated
+    into the offset amplitudes.
+
+    Args:
+        data_index (int):  The detector index to process.
+        det_data (array, double):  The data for all detectors.
+        flag_index (int):  The detector index in the flag data.  A negative value
+            indicates no flags.
+        flag_data (array, uint8):  The flags for all detectors.
+        flag_mask (uint8):  The flag mask to apply.
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int64):  The first amplitude for this detector.
+        n_amp_views (array, int):  The number of amplitudes for each interval.
+        amplitudes (array, double): The amplitude data.
+        intervals (array, Interval):  The intervals to process.
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return offset_project_signal(
+        data_index,
+        det_data,
+        flag_index,
+        flag_data,
+        flag_mask,
+        step_length,
+        amp_offset,
+        n_amp_views,
+        amplitudes,
+        intervals,
+        impl=ImplementationType.COMPILED,
+        use_accel=use_accel,
+    )
+
+
+@kernel(impl=ImplementationType.DEFAULT)
+def offset_apply_diag_precond(
+    offset_var,
+    amplitudes_in,
+    amplitudes_out,
+    use_accel=False,
+):
+    """
+    Args:
+        offset_var (array, double):  The variance weight to apply to each amplitude.
+        amplitudes_in (array, double):  Input amplitude data
+        amplitudes_out (array, double):  Output amplitude data
+        use_accel (bool):  Whether to use the accelerator for this call (if supported).
+
+    Returns:
+        None
+
+    """
+    return offset_apply_diag_precond(
+        offset_var,
+        amplitudes_in,
+        amplitudes_out,
+        impl=ImplementationType.COMPILED,
+        use_accel=False,
+    )
+
+
+@kernel(impl=ImplementationType.COMPILED, name="offset_add_to_signal")
+def offset_add_to_signal_compiled(*args, use_accel=False):
+    return libtoast_offset_add_to_signal(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="offset_project_signal")
+def offset_project_signal_compiled(*args, use_accel=False):
+    return libtoast_offset_project_signal(*args, use_accel)
+
+
+@kernel(impl=ImplementationType.COMPILED, name="offset_apply_diag_precond")
+def offset_apply_diag_precond_compiled(*args, use_accel=False):
+    return libtoast_offset_apply_diag_precond(*args, use_accel)
diff --git a/src/toast/templates/offset/kernels_jax.py b/src/toast/templates/offset/kernels_jax.py
new file mode 100644
index 000000000..1016210fd
--- /dev/null
+++ b/src/toast/templates/offset/kernels_jax.py
@@ -0,0 +1,358 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+from ...jax.intervals import INTERVALS_JAX, JaxIntervals
+from ...jax.mutableArray import MutableJaxArray
+from ...utils import Logger
+
+
+def offset_add_to_signal_intervals(
+    step_length,
+    amplitudes,
+    data_index,
+    det_data,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+    offset_starts,
+    offset_ends,
+    offsets_max_length,
+):
+    """
+    Each amplitude value is accumulated to `step_length` number of samples.  The
+    final offset will be at least this many samples, but may be more if the step
+    size does not evenly divide into the number of samples.
+
+    Process all the intervals as a block.
+
+    Args:
+        step_length (int64):  The minimum number of samples for each offset.
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+        offset_starts (array, int): size n_view
+        offset_ends (array, int): size n_view
+        offsets_max_length (int): maximum length of an interval
+
+    Returns:
+        det_data
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"offset_add_to_signal: jit-compiling.")
+
+    # split data to separate the final amplitude from the rest
+    # as it is the only one that does not have step_length samples
+    nb_intervals = interval_starts.size
+    nb_amplitudes = offsets_max_length
+
+    # computes interval data
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    offsets = JaxIntervals(
+        offset_starts, offset_ends + 1, offsets_max_length
+    )  # end+1 as the interval is inclusive
+    amplitudes_interval = JaxIntervals.get(amplitudes, offsets)  # amplitudes[offsets]
+    data_interval = JaxIntervals.get(
+        det_data, (data_index, intervals)
+    )  # det_data[data_index, intervals]
+
+    # All but the last amplitude have step_length samples.
+    data_first = data_interval[:, : (nb_amplitudes - 1) * step_length]
+    data_first = jnp.reshape(data_first, newshape=(nb_intervals, -1, step_length))
+    new_data_first = data_first + amplitudes_interval[:, :-1, jnp.newaxis]
+    # data_first[:] += amplitudes[:-1, jnp.newaxis]
+    data_interval = data_interval.at[:, : (nb_amplitudes - 1) * step_length].set(
+        new_data_first.reshape((nb_intervals, -1))
+    )
+
+    # Now handle the final amplitude.
+    # data_last = data[(nb_amplitudes - 1) * step_length:]
+    # data_last[:] += amplitudes[-1]
+    data_interval = data_interval.at[:, (nb_amplitudes - 1) * step_length :].add(
+        amplitudes_interval[:, -1]
+    )
+
+    # updates det_data and returns
+    # det_data[data_index, intervals] = data_interval
+    det_data = JaxIntervals.set(det_data, (data_index, intervals), data_interval)
+    return det_data
+
+
+# jit compilation
+offset_add_to_signal_intervals = jax.jit(
+    offset_add_to_signal_intervals,
+    static_argnames=["step_length", "intervals_max_length", "offsets_max_length"],
+    donate_argnums=[3],
+)  # det_data
+
+
+@kernel(impl=ImplementationType.JAX, name="offset_add_to_signal")
+def offset_add_to_signal_jax(
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    data_index,
+    det_data,
+    intervals,
+    use_accel,
+):
+    """
+    Accumulate offset amplitudes to timestream data.
+    Each amplitude value is accumulated to `step_length` number of samples.
+
+    Args:
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int): starting offset
+        n_amp_views (array, int): subsequent offsets (size n_view)
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        intervals (array, Interval): size n_view
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in det_data).
+    """
+    # prepare inputs
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    det_data_input = MutableJaxArray.to_array(det_data)
+    amplitudes = MutableJaxArray.to_array(amplitudes)
+
+    # computes offsets
+    # a cumulative sums of the n_amp_views starting at amp_offset
+    offsets = np.roll(n_amp_views, 1)
+    offsets[0] = amp_offset
+    offsets = np.cumsum(offsets)
+    # prepare offsets intervals
+    offsets_start = offsets + intervals.first // step_length
+    offsets_end = offsets + intervals.last // step_length
+    offsets_max_length = int(np.max(1 + offsets_end - offsets_start))
+
+    # run computation
+    det_data[:] = offset_add_to_signal_intervals(
+        step_length,
+        amplitudes,
+        data_index,
+        det_data_input,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+        offsets_start,
+        offsets_end,
+        offsets_max_length,
+    )
+
+
+def offset_project_signal_intervals(
+    data_index,
+    det_data,
+    use_flag,
+    flag_index,
+    flag_data,
+    flag_mask,
+    step_length,
+    amplitudes,
+    interval_starts,
+    interval_ends,
+    intervals_max_length,
+    offset_starts,
+    offset_ends,
+    offsets_max_length,
+):
+    """
+    Chunks of `step_length` number of samples are accumulated into the offset
+    amplitudes.  If step_length does not evenly divide into the total number of
+    samples, the final amplitude will be extended to include the remainder.
+
+    Process all the intervals as a block.
+
+    Args:
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        use_flag (bool): should we use flags
+        flag_index (int), strictly negative in the absence of a detcetor flag
+        flag_data (array, bool) size n_all_det*n_samp
+        flag_mask (int),
+        step_length (int64):  The minimum number of samples for each offset.
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        interval_starts (array, int): size n_view
+        interval_ends (array, int): size n_view
+        intervals_max_length (int): maximum length of an interval
+        offset_starts (array, int): size n_view
+        offset_ends (array, int): size n_view
+        offsets_max_length (int): maximum length of an interval
+
+    Returns:
+        amplitudes
+    """
+    # debugging information
+    log = Logger.get()
+    log.debug(f"offset_project_signal: jit-compiling.")
+
+    # gets interval information
+    nb_amplitudes = offsets_max_length
+    nb_intervals = interval_starts.size
+
+    # computes interval data
+    intervals = JaxIntervals(
+        interval_starts, interval_ends + 1, intervals_max_length
+    )  # end+1 as the interval is inclusive
+    offsets = JaxIntervals(
+        offset_starts, offset_ends + 1, offsets_max_length
+    )  # end+1 as the interval is inclusive
+    amplitudes_interval = JaxIntervals.get(amplitudes, offsets)  # amplitudes[offsets]
+    det_data_interval = JaxIntervals.get(
+        det_data, (data_index, intervals), padding_value=0.0
+    )  # det_data[data_index,intervals]
+
+    # skip flagged samples
+    if use_flag:
+        flags_interval = JaxIntervals.get(
+            flag_data, (flag_index, intervals)
+        )  # flag_data[flag_index,intervals]
+        flagged = (flags_interval & flag_mask) != 0
+        det_data_interval = jnp.where(flagged, 0.0, det_data_interval)
+
+    # split data to separate the final amplitude from the rest
+    # as it is the only one that does not have step_length samples
+    data_first = det_data_interval[:, : (nb_amplitudes - 1) * step_length]
+    data_first = jnp.reshape(data_first, newshape=(nb_intervals, -1, step_length))
+    data_last = det_data_interval[:, (nb_amplitudes - 1) * step_length :]
+
+    # All but the last amplitude have step_length samples.
+    # amplitudes[:-1] += np.sum(data_first, axis=1)
+    amplitudes_interval = amplitudes_interval.at[:, :-1].add(
+        jnp.sum(data_first, axis=-1)
+    )
+
+    # Now handle the final amplitude.
+    # amplitudes[-1] += np.sum(data_last)
+    amplitudes_interval = amplitudes_interval.at[:, -1].add(jnp.sum(data_last, axis=-1))
+
+    # updates amplitudes and returns
+    amplitudes = JaxIntervals.set(
+        amplitudes, offsets, amplitudes_interval
+    )  # amplitudes[offsets] = amplitudes_interval
+    return amplitudes
+
+
+# jit compilation
+offset_project_signal_intervals = jax.jit(
+    offset_project_signal_intervals,
+    static_argnames=[
+        "use_flag",
+        "flag_mask",
+        "step_length",
+        "intervals_max_length",
+        "offsets_max_length",
+    ],
+    donate_argnums=[7],
+)  # donate amplitude
+
+
+@kernel(impl=ImplementationType.JAX, name="offset_project_signal")
+def offset_project_signal_jax(
+    data_index,
+    det_data,
+    flag_index,
+    flag_data,
+    flag_mask,
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    intervals,
+    use_accel,
+):
+    """
+    Accumulate timestream data into offset amplitudes.
+    Chunks of `step_length` number of samples are accumulated into the offset amplitudes.
+
+    Args:
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        flag_index (int), strictly negative in the absence of a detcetor flag
+        flag_data (array, bool) size n_all_det*n_samp
+        flag_mask (int),
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int)
+        n_amp_views (array, int): size n_view
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        intervals (array, Interval): size n_view
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in amplitudes).
+    """
+    # prepare inputs
+    intervals_max_length = INTERVALS_JAX.compute_max_intervals_length(intervals)
+    use_flag = flag_index >= 0
+    det_data = MutableJaxArray.to_array(det_data)
+    flag_data = MutableJaxArray.to_array(flag_data)
+    amplitudes_input = MutableJaxArray.to_array(amplitudes)
+
+    # computes offsets
+    # a cumulative sums of the n_amp_views starting at amp_offset
+    offsets = np.roll(n_amp_views, 1)
+    offsets[0] = amp_offset
+    offsets = np.cumsum(offsets)
+    # prepare offsets intervals
+    offsets_start = offsets + intervals.first // step_length
+    offsets_end = offsets + intervals.last // step_length
+    offsets_max_length = int(np.max(1 + offsets_end - offsets_start))
+
+    # run computation
+    amplitudes[:] = offset_project_signal_intervals(
+        data_index,
+        det_data,
+        use_flag,
+        flag_index,
+        flag_data,
+        flag_mask,
+        step_length,
+        amplitudes_input,
+        intervals.first,
+        intervals.last,
+        intervals_max_length,
+        offsets_start,
+        offsets_end,
+        offsets_max_length,
+    )
+
+
+@kernel(impl=ImplementationType.JAX, name="offset_apply_diag_precond")
+def offset_apply_diag_precond_jax(offset_var, amplitudes_in, amplitudes_out, use_accel):
+    """
+    NOTE this does not use JAX as there is too little computation
+
+    Args:
+        offset_var (array, double): size n_amp
+        amplitudes_in (array, double): size n_amp
+        amplitudes_out (array, double): size n_amp
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in amplitudes_out).
+    """
+    # cast the data if needed
+    offset_var = MutableJaxArray.to_array(offset_var)
+    amplitudes_in = MutableJaxArray.to_array(amplitudes_in)
+
+    # runs the computation
+    amplitudes_out[:] = amplitudes_in * offset_var
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("template_offset"); toast.tests.run("ops_mapmaker_solve"); toast.tests.run("ops_mapmaker")'
diff --git a/src/toast/templates/offset/kernels_numpy.py b/src/toast/templates/offset/kernels_numpy.py
new file mode 100644
index 000000000..bf3d877e3
--- /dev/null
+++ b/src/toast/templates/offset/kernels_numpy.py
@@ -0,0 +1,121 @@
+# Copyright (c) 2015-2023 by the parties listed in the AUTHORS file.
+# All rights reserved.  Use of this source code is governed by
+# a BSD-style license that can be found in the LICENSE file.
+
+import numpy as np
+
+from ...accelerator import ImplementationType, kernel
+
+
+@kernel(impl=ImplementationType.NUMPY, name="offset_add_to_signal")
+def offset_add_to_signal_numpy(
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    data_index,
+    det_data,
+    intervals,
+    use_accel,
+):
+    """
+    Accumulate offset amplitudes to timestream data.
+    Each amplitude value is accumulated to `step_length` number of samples.
+
+    Args:
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int)
+        n_amp_views (array, int): size n_view
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        intervals (array, Interval): size n_view
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in det_data).
+    """
+    offset = amp_offset
+    for interval, view_offset in zip(intervals, n_amp_views):
+        samples = slice(interval.first, interval.last + 1, 1)
+        sampidx = np.arange(interval.first, interval.last + 1, dtype=np.int64)
+        amp_vals = np.array([amplitudes[offset + x] for x in (sampidx // step_length)])
+        det_data[data_index, samples] += amp_vals
+        offset += view_offset
+
+
+@kernel(impl=ImplementationType.NUMPY, name="offset_project_signal")
+def offset_project_signal_numpy(
+    data_index,
+    det_data,
+    flag_index,
+    flag_data,
+    flag_mask,
+    step_length,
+    amp_offset,
+    n_amp_views,
+    amplitudes,
+    intervals,
+    use_accel,
+):
+    """
+    Accumulate timestream data into offset amplitudes.
+    Chunks of `step_length` number of samples are accumulated into the offset amplitudes.
+
+    Args:
+        data_index (int)
+        det_data (array, double): The float64 timestream values (size n_all_det*n_samp).
+        flag_index (int), strictly negative in the absence of a detcetor flag
+        flag_data (array, bool) size n_all_det*n_samp
+        flag_mask (int),
+        step_length (int64):  The minimum number of samples for each offset.
+        amp_offset (int)
+        n_amp_views (array, int): size n_view
+        amplitudes (array, double): The float64 amplitude values (size n_amp)
+        intervals (array, Interval): size n_view
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in amplitudes).
+    """
+    offset = amp_offset
+    for interval, view_offset in zip(intervals, n_amp_views):
+        samples = slice(interval.first, interval.last + 1, 1)
+        ampidx = (
+            offset
+            + np.arange(interval.first, interval.last + 1, dtype=np.int64)
+            // step_length
+        )
+        ddata = det_data[data_index][samples]
+        # skip sample if it is flagged
+        if flag_index >= 0:
+            # We have detector flags
+            ddata = np.array(
+                ((flag_data[flag_index] & flag_mask) == 0), dtype=np.float64
+            )
+            ddata *= det_data[data_index][samples]
+        # updates amplitude
+        # using np.add to insure atomicity
+        np.add.at(amplitudes, ampidx, ddata)
+        offset += view_offset
+
+
+@kernel(impl=ImplementationType.NUMPY, name="offset_apply_diag_precond")
+def offset_apply_diag_precond_numpy(
+    offset_var, amplitudes_in, amplitudes_out, use_accel
+):
+    """
+    Args:
+        offset_var (array, double): size n_amp
+        amplitudes_in (array, double): size n_amp
+        amplitudes_out (array, double): size n_amp
+        use_accel (bool): should we use the accelerator
+
+    Returns:
+        None (the result is put in amplitudes_out).
+    """
+    amplitudes_out[:] = amplitudes_in * offset_var
+
+
+# To test:
+# python -c 'import toast.tests; toast.tests.run("template_offset"); toast.tests.run("ops_mapmaker_solve"); toast.tests.run("ops_mapmaker")'
diff --git a/src/toast/templates/offset.py b/src/toast/templates/offset/offset.py
similarity index 83%
rename from src/toast/templates/offset.py
rename to src/toast/templates/offset/offset.py
index 74c08753a..fe1d7b962 100644
--- a/src/toast/templates/offset.py
+++ b/src/toast/templates/offset/offset.py
@@ -10,19 +10,19 @@
 import traitlets
 from astropy import units as u
 
-from .._libtoast import (
-    template_offset_add_to_signal,
-    template_offset_apply_diag_precond,
-    template_offset_project_signal,
+from ...accelerator import ImplementationType
+from ...mpi import MPI
+from ...observation import default_values as defaults
+from ...timing import function_timer
+from ...traits import Bool, Float, Int, Quantity, Unicode, UseEnum, trait_docs
+from ...utils import AlignedF64, Logger, rate_from_times
+from ..amplitudes import Amplitudes
+from ..template import Template
+from .kernels import (
+    offset_add_to_signal,
+    offset_apply_diag_precond,
+    offset_project_signal,
 )
-from ..data import Data
-from ..mpi import MPI
-from ..observation import default_values as defaults
-from ..timing import function_timer
-from ..traits import Bool, Float, Instance, Int, Quantity, Unicode, trait_docs
-from ..utils import AlignedF64, Logger, rate_from_times
-from .amplitudes import Amplitudes
-from .template import Template
 
 
 @trait_docs
@@ -68,8 +68,6 @@ class Offset(Template):
 
     precond_width = Int(20, help="Preconditioner width in terms of offsets / baselines")
 
-    use_python = Bool(False, help="If True, use python implementation")
-
     @traitlets.validate("precond_width")
     def _check_precond_width(self, proposal):
         w = proposal["value"]
@@ -105,9 +103,6 @@ def _initialize(self, new_data):
         # Clear any C-allocated buffers from previous uses.
         self.clear()
 
-        if self.use_python and self.use_accel:
-            raise RuntimeError("Cannot use accelerator with pure python implementation")
-
         # Compute the step boundaries for every observation and the number of
         # amplitude values on this process.  Every process only stores amplitudes
         # for its locally assigned data.
@@ -495,14 +490,17 @@ def _step_length(self, stime, rate):
         return int(stime * rate + 0.5)
 
     @function_timer
-    def _add_to_signal(self, detector, amplitudes):
+    def _add_to_signal(self, detector, amplitudes, use_accel=False, **kwargs):
         log = Logger.get()
+
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         amp_offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
             if detector not in ob.local_detectors:
                 continue
             det_indx = ob.detdata[self.det_data].indices([detector])
-
             # The step length for this observation
             step_length = self._step_length(
                 self.step_time.to_value(u.second), self._obs_rate[iob]
@@ -511,32 +509,27 @@ def _add_to_signal(self, detector, amplitudes):
             # The number of amplitudes in each view
             n_amp_views = self._obs_views[iob]
 
-            if self.use_python:
-                self._py_add_to_signal(
-                    step_length,
-                    amp_offset,
-                    n_amp_views,
-                    amplitudes.local,
-                    det_indx,
-                    ob.detdata[self.det_data].data,
-                    ob.intervals[self.view].data,
-                )
-            else:
-                template_offset_add_to_signal(
-                    step_length,
-                    amp_offset,
-                    n_amp_views,
-                    amplitudes.local,
-                    det_indx[0],
-                    ob.detdata[self.det_data].data,
-                    ob.intervals[self.view].data,
-                    self.use_accel,
-                )
+            offset_add_to_signal(
+                step_length,
+                amp_offset,
+                n_amp_views,
+                amplitudes.local,
+                det_indx[0],
+                ob.detdata[self.det_data].data,
+                ob.intervals[self.view].data,
+                impl=implementation,
+                use_accel=use_accel,
+            )
+
             amp_offset += np.sum(n_amp_views)
 
     @function_timer
-    def _project_signal(self, detector, amplitudes):
+    def _project_signal(self, detector, amplitudes, use_accel=False, **kwargs):
         log = Logger.get()
+
+        # Kernel selection
+        implementation = self.select_kernels(use_accel=use_accel)
+
         amp_offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
             if detector not in ob.local_detectors:
@@ -556,41 +549,29 @@ def _project_signal(self, detector, amplitudes):
             # The number of amplitudes in each view
             n_amp_views = self._obs_views[iob]
 
-            if self.use_python:
-                self._py_project_signal(
-                    det_indx,
-                    ob.detdata[self.det_data].data,
-                    flag_indx,
-                    flag_data,
-                    self.det_flag_mask,
-                    step_length,
-                    amp_offset,
-                    n_amp_views,
-                    amplitudes.local,
-                    ob.intervals[self.view].data,
-                )
-            else:
-                template_offset_project_signal(
-                    det_indx[0],
-                    ob.detdata[self.det_data].data,
-                    flag_indx[0],
-                    flag_data,
-                    self.det_flag_mask,
-                    step_length,
-                    amp_offset,
-                    n_amp_views,
-                    amplitudes.local,
-                    ob.intervals[self.view].data,
-                    self.use_accel,
-                )
+            offset_project_signal(
+                det_indx[0],
+                ob.detdata[self.det_data].data,
+                flag_indx[0],
+                flag_data,
+                self.det_flag_mask,
+                step_length,
+                amp_offset,
+                n_amp_views,
+                amplitudes.local,
+                ob.intervals[self.view].data,
+                impl=implementation,
+                use_accel=use_accel,
+            )
+
             amp_offset += np.sum(n_amp_views)
 
     @function_timer
-    def _add_prior(self, amplitudes_in, amplitudes_out):
+    def _add_prior(self, amplitudes_in, amplitudes_out, use_accel=False, **kwargs):
         if not self.use_noise_prior:
             # Not using the noise prior term, nothing to accumulate to output.
             return
-        if self.use_accel:
+        if use_accel:
             raise NotImplementedError(
                 "offset template add_prior on accelerator not implemented"
             )
@@ -610,9 +591,9 @@ def _add_prior(self, amplitudes_in, amplitudes_out):
                     offset += n_amp_view
 
     @function_timer
-    def _apply_precond(self, amplitudes_in, amplitudes_out):
+    def _apply_precond(self, amplitudes_in, amplitudes_out, use_accel=False, **kwargs):
         if self.use_noise_prior:
-            if self.use_accel:
+            if use_accel:
                 raise NotImplementedError(
                     "offset template precond on accelerator not implemented"
                 )
@@ -659,76 +640,26 @@ def _apply_precond(self, amplitudes_in, amplitudes_out):
         else:
             # Since we do not have a noise filter term in our LHS, our diagonal
             # preconditioner is just the application of offset variance.
-            if self.use_python:
-                self._py_apply_diag_precond(
-                    self._offsetvar,
-                    amplitudes_in.local,
-                    amplitudes_out.local,
-                )
-            else:
-                template_offset_apply_diag_precond(
-                    self._offsetvar,
-                    amplitudes_in.local,
-                    amplitudes_out.local,
-                    self.use_accel,
-                )
+
+            # Kernel selection
+            implementation = self.select_kernels(use_accel=use_accel)
+
+            offset_apply_diag_precond(
+                self._offsetvar,
+                amplitudes_in.local,
+                amplitudes_out.local,
+                impl=implementation,
+                use_accel=use_accel,
+            )
         return
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+            ImplementationType.COMPILED,
+            ImplementationType.NUMPY,
+            ImplementationType.JAX,
+        ]
+
     def _supports_accel(self):
         return True
-
-    def _py_add_to_signal(
-        self,
-        step_length,
-        amp_offset,
-        n_amp_views,
-        amplitudes,
-        data_index,
-        det_data,
-        intr_data,
-    ):
-        """Internal python implementation for comparison testing."""
-        offset = amp_offset
-        for ivw, vw in enumerate(intr_data):
-            samples = slice(vw.first, vw.last + 1, 1)
-            sampidx = np.arange(vw.first, vw.last + 1, dtype=np.int64)
-            amp_vals = np.array(
-                [amplitudes[offset + x] for x in (sampidx // step_length)]
-            )
-            det_data[data_index[0], samples] += amp_vals
-            offset += n_amp_views[ivw]
-
-    def _py_project_signal(
-        self,
-        data_index,
-        det_data,
-        flag_index,
-        flag_data,
-        flag_mask,
-        step_length,
-        amp_offset,
-        n_amp_views,
-        amplitudes,
-        intr_data,
-    ):
-        """Internal python implementation for comparison testing."""
-        offset = amp_offset
-        for ivw, vw in enumerate(intr_data):
-            samples = slice(vw.first, vw.last + 1, 1)
-            ampidx = (
-                offset + np.arange(vw.first, vw.last + 1, dtype=np.int64) // step_length
-            )
-            ddata = det_data[data_index[0]][samples]
-            if flag_index[0] >= 0:
-                # We have detector flags
-                ddata = np.array(
-                    ((flag_data[flag_index[0]] & flag_mask) == 0), dtype=np.float64
-                )
-                ddata *= det_data[data_index[0]][samples]
-            np.add.at(amplitudes, ampidx, ddata)
-            offset += n_amp_views[ivw]
-
-    def _py_apply_diag_precond(self, offset_var, amp_in, amp_out):
-        """Internal python implementation for comparison testing."""
-        amp_out[:] = amp_in
-        amp_out *= offset_var
diff --git a/src/toast/templates/subharmonic.py b/src/toast/templates/subharmonic.py
index 393e17bde..946117a31 100644
--- a/src/toast/templates/subharmonic.py
+++ b/src/toast/templates/subharmonic.py
@@ -169,7 +169,7 @@ def _zeros(self):
         # z.local_flags[:] = np.where(self._amp_flags, 1, 0)
         return z
 
-    def _add_to_signal(self, detector, amplitudes):
+    def _add_to_signal(self, detector, amplitudes, **kwargs):
         norder = self.order + 1
         offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
@@ -181,7 +181,7 @@ def _add_to_signal(self, detector, amplitudes):
                     vw[detector] += self._templates[iob][ivw][order] * amp_view[order]
                 offset += norder
 
-    def _project_signal(self, detector, amplitudes):
+    def _project_signal(self, detector, amplitudes, **kwargs):
         norder = self.order + 1
         offset = self._det_start[detector]
         for iob, ob in enumerate(self.data.obs):
@@ -193,16 +193,16 @@ def _project_signal(self, detector, amplitudes):
                     amp_view[order] = np.dot(vw[detector], template)
                 offset += norder
 
-    def _add_prior(self, amplitudes_in, amplitudes_out):
+    def _add_prior(self, amplitudes_in, amplitudes_out, **kwargs):
         # No prior for this template, nothing to accumulate to output.
         return
 
-    def _apply_precond(self, amplitudes_in, amplitudes_out):
+    def _apply_precond(self, amplitudes_in, amplitudes_out, **kwargs):
         norder = self.order + 1
         for det in self._all_dets:
-            offset = self._det_start[detector]
+            offset = self._det_start[det]
             for iob, ob in enumerate(self.data.obs):
-                if detector not in ob.local_detectors:
+                if det not in ob.local_detectors:
                     continue
                 views = ob.view[self.view]
                 for ivw, vw in enumerate(views):
diff --git a/src/toast/templates/template.py b/src/toast/templates/template.py
index 23a859c9f..a0d523b1b 100644
--- a/src/toast/templates/template.py
+++ b/src/toast/templates/template.py
@@ -62,8 +62,6 @@ class Template(TraitConfig):
         defaults.det_mask_invalid, help="Bit mask value for solver flags"
     )
 
-    use_accel = Bool(False, help="If True, use the accelerator")
-
     @traitlets.validate("data")
     def _check_data(self, proposal):
         dat = proposal["value"]
@@ -86,12 +84,16 @@ def _initialize(self, new_data):
         # computing the number of amplitudes) whenever the data changes.
         raise NotImplementedError("Derived class must implement _initialize()")
 
-    def _check_enabled(self):
+    def _check_enabled(self, use_accel=False):
         if self.data is None:
             raise RuntimeError(
                 "You must set the data trait before calling template methods"
             )
         if self.enabled:
+            if use_accel and not self.supports_accel():
+                msg = f"Template {self.name} does not support accelerator, "
+                msg += "cannot specify use_accel=True"
+                raise RuntimeError(msg)
             return True
         else:
             log = Logger.get()
@@ -134,11 +136,11 @@ def zeros(self):
         if self._check_enabled():
             return self._zeros()
 
-    def _add_to_signal(self, detector, amplitudes):
+    def _add_to_signal(self, detector, amplitudes, **kwargs):
         raise NotImplementedError("Derived class must implement _add_to_signal()")
 
     @function_timer_stackskip
-    def add_to_signal(self, detector, amplitudes):
+    def add_to_signal(self, detector, amplitudes, use_accel=False, **kwargs):
         """Accumulate the projected amplitudes to a timestream.
 
         This performs the operation:
@@ -156,14 +158,14 @@ def add_to_signal(self, detector, amplitudes):
             None
 
         """
-        if self._check_enabled():
-            return self._add_to_signal(detector, amplitudes)
+        if self._check_enabled(use_accel=use_accel):
+            self._add_to_signal(detector, amplitudes, use_accel=use_accel, **kwargs)
 
-    def _project_signal(self, detector, amplitudes):
+    def _project_signal(self, detector, amplitudes, **kwargs):
         raise NotImplementedError("Derived class must implement _project_signal()")
 
     @function_timer_stackskip
-    def project_signal(self, detector, amplitudes):
+    def project_signal(self, detector, amplitudes, use_accel=False, **kwargs):
         """Project a timestream into template amplitudes.
 
         This performs:
@@ -181,15 +183,15 @@ def project_signal(self, detector, amplitudes):
             None
 
         """
-        if self._check_enabled():
-            return self._project_signal(detector, amplitudes)
+        if self._check_enabled(use_accel=use_accel):
+            self._project_signal(detector, amplitudes, use_accel=use_accel, **kwargs)
 
-    def _add_prior(self, amplitudes_in, amplitudes_out):
+    def _add_prior(self, amplitudes_in, amplitudes_out, **kwargs):
         # Not all Templates implement the prior
         return
 
     @function_timer_stackskip
-    def add_prior(self, amplitudes_in, amplitudes_out):
+    def add_prior(self, amplitudes_in, amplitudes_out, use_accel=False, **kwargs):
         """Apply the inverse amplitude covariance as a prior.
 
         This performs:
@@ -205,14 +207,14 @@ def add_prior(self, amplitudes_in, amplitudes_out):
             None
 
         """
-        if self._check_enabled():
-            return self._add_prior(amplitudes_in, amplitudes_out)
+        if self._check_enabled(use_accel=use_accel):
+            self._add_prior(amplitudes_in, amplitudes_out, use_accel=use_accel, **kwargs)
 
-    def _apply_precond(self, amplitudes_in, amplitudes_out):
+    def _apply_precond(self, amplitudes_in, amplitudes_out, **kwargs):
         raise NotImplementedError("Derived class must implement _apply_precond()")
 
     @function_timer_stackskip
-    def apply_precond(self, amplitudes_in, amplitudes_out):
+    def apply_precond(self, amplitudes_in, amplitudes_out, use_accel=False, **kwargs):
         """Apply the template preconditioner.
 
         Formally, the preconditioner "M" is an approximation to the "design matrix"
@@ -230,20 +232,8 @@ def apply_precond(self, amplitudes_in, amplitudes_out):
             None
 
         """
-        if self._check_enabled():
-            return self._apply_precond(amplitudes_in, amplitudes_out)
-
-    def _supports_accel(self):
-        return False
-
-    def supports_accel(self):
-        """Query whether the template supports OpenACC
-
-        Returns:
-            (bool):  True if the template can use OpenACC, else False.
-
-        """
-        return self._supports_accel()
+        if self._check_enabled(use_accel=use_accel):
+            self._apply_precond(amplitudes_in, amplitudes_out, use_accel=use_accel, **kwargs)
 
     @classmethod
     def get_class_config_path(cls):
diff --git a/src/toast/traits.py b/src/toast/traits.py
index 8ca972c76..4aba0b684 100644
--- a/src/toast/traits.py
+++ b/src/toast/traits.py
@@ -24,10 +24,12 @@
     Tuple,
     Undefined,
     Unicode,
+    UseEnum,
     signature_has_traits,
 )
 
-from .utils import import_from_name, object_fullname
+from .accelerator import ImplementationType, use_accel_jax, use_accel_omp
+from .utils import import_from_name, object_fullname, Logger
 
 
 def trait_string_to_scalar(val):
@@ -134,6 +136,7 @@ def trait_scalar_to_string(val):
 # Scalar base types
 
 Bool.py_type = lambda self: bool
+UseEnum.py_type = lambda self: self.enum_class
 Float.py_type = lambda self: float
 Int.py_type = lambda self: int
 Unicode.py_type = lambda self: str
@@ -158,6 +161,7 @@ def _get_conf(self, obj=None):
 
 
 Bool.get_conf = _create_scalar_trait_get_conf("bool")
+UseEnum.get_conf = _create_scalar_trait_get_conf("enum")
 Float.get_conf = _create_scalar_trait_get_conf("float")
 Int.get_conf = _create_scalar_trait_get_conf("int")
 Unicode.get_conf = _create_scalar_trait_get_conf("str")
@@ -459,7 +463,14 @@ class constructor.
 
     enabled = Bool(True, help="If True, this class instance is marked as enabled")
 
+    kernel_implementation = UseEnum(
+        ImplementationType,
+        default_value=ImplementationType.DEFAULT,
+        help="Which kernel implementation to use (DEFAULT, COMPILED, NUMPY, JAX).",
+    )
+
     def __init__(self, **kwargs):
+        self._accel_stack = list()
         super().__init__(**kwargs)
         if self.name is None:
             self.name = self.__class__.__qualname__
@@ -471,6 +482,58 @@ def __repr__(self):
         val += "\n>"
         return val
 
+    def _implementations(self):
+        return [
+            ImplementationType.DEFAULT,
+        ]
+
+    def implementations(self):
+        """Query which kernel implementations are supported.
+
+        Returns:
+            (list):  List of implementations.
+
+        """
+        return self._implementations()
+
+    def _supports_accel(self):
+        return False
+
+    def supports_accel(self):
+        """Query whether the operator supports accelerator kernels
+
+        Returns:
+            (bool):  True if the operator can use accelerators, else False.
+
+        """
+        return self._supports_accel()
+
+    def select_kernels(self, use_accel=False):
+        """Return the currently selected kernel implementation.
+
+        This returns the kernel implementation that should be used
+
+        Returns:
+            (ImplementationType):  The implementation type.
+
+        """
+        impls = self.implementations()
+        if use_accel:
+            if use_accel_jax:
+                if ImplementationType.JAX not in impls:
+                    msg = f"JAX accelerator use is enabled, "
+                    msg += f"but not supported by {self.name}"
+                    raise RuntimeError(msg)
+                return ImplementationType.JAX
+            else:
+                if ImplementationType.COMPILED not in impls:
+                    msg = f"OpenMP accelerator use is enabled, "
+                    msg += f"but not supported by {self.name}"
+                    raise RuntimeError(msg)
+                return ImplementationType.COMPILED
+        else:
+            return ImplementationType.DEFAULT
+
     def __eq__(self, other):
         if len(self.traits()) != len(other.traits()):
             # print(


Test patch:
diff --git a/src/toast/tests/accelerator.py b/src/toast/tests/accelerator.py
index 9ad87835c..56b19ffed 100644
--- a/src/toast/tests/accelerator.py
+++ b/src/toast/tests/accelerator.py
@@ -11,12 +11,14 @@
 from .. import ops
 from .._libtoast import test_accel_op_array, test_accel_op_buffer
 from ..accelerator import (
+    ImplementationType,
     accel_data_create,
     accel_data_delete,
     accel_data_present,
     accel_data_update_device,
     accel_data_update_host,
     accel_enabled,
+    kernel,
     use_accel_jax,
     use_accel_omp,
 )
@@ -90,6 +92,31 @@ def setUp(self):
             "u8": np.uint8,
         }
 
+    def test_kernel_registry(self):
+        @kernel(ImplementationType.DEFAULT)
+        def my_kernel(foo, use_accel=False):
+            print(f"DEFAULT (accel={use_accel}):  {foo}")
+
+        @kernel(ImplementationType.COMPILED, name="my_kernel")
+        def super_compiled(foo, use_accel=False):
+            print(f"COMPILED (accel={use_accel}):  {foo}")
+
+        @kernel(ImplementationType.NUMPY, name="my_kernel")
+        def super_numpy(foo, use_accel=False):
+            print(f"NUMPY (accel={use_accel}):  {foo}")
+
+        @kernel(ImplementationType.JAX, name="my_kernel")
+        def awesome_jax(foo, use_accel=False):
+            print(f"JAX (accel={use_accel}):  {foo}")
+
+        bar = "yes"
+        my_kernel(bar, impl=ImplementationType.DEFAULT)
+        my_kernel(bar, impl=ImplementationType.NUMPY)
+        my_kernel(bar, impl=ImplementationType.COMPILED)
+        my_kernel(bar, impl=ImplementationType.COMPILED, use_accel=True)
+        my_kernel(bar, impl=ImplementationType.JAX)
+        my_kernel(bar, impl=ImplementationType.JAX, use_accel=True)
+
     def test_memory(self):
         if not (use_accel_omp or use_accel_jax):
             if self.rank == 0:
@@ -107,7 +134,7 @@ def test_memory(self):
 
         # Copy to device
         for tname, buffer in data.items():
-            accel_data_create(buffer)
+            accel_data_create(buffer)  # TODO buffer = acel... in the jax case?
             accel_data_update_device(buffer)
 
         # Check that it is present
diff --git a/src/toast/tests/config.py b/src/toast/tests/config.py
index 19f9f89c0..38a159aa3 100644
--- a/src/toast/tests/config.py
+++ b/src/toast/tests/config.py
@@ -247,10 +247,6 @@ def _requires(self):
     def _provides(self):
         return dict()
 
-    def _accelerators(self):
-        return list()
-
-
 class ConfigTest(MPITestCase):
     def setUp(self):
         fixture_name = os.path.splitext(os.path.basename(__file__))[0]
diff --git a/src/toast/tests/covariance.py b/src/toast/tests/covariance.py
index d8dd657f7..8878dd801 100644
--- a/src/toast/tests/covariance.py
+++ b/src/toast/tests/covariance.py
@@ -149,6 +149,7 @@ def test_multiply(self):
                             npp.data[sm, spix, elem], chk, rtol=1e-5, atol=1e-8
                         ):
                             failed = True
+                            print(f"{npp.data[sm, spix, elem]} differs from {chk}")
             if comm is not None:
                 failed = comm.allreduce(failed, op=MPI.LOR)
             self.assertFalse(failed)
diff --git a/src/toast/tests/ops_mapmaker.py b/src/toast/tests/ops_mapmaker.py
index 8ca6d85d8..63764cd9d 100644
--- a/src/toast/tests/ops_mapmaker.py
+++ b/src/toast/tests/ops_mapmaker.py
@@ -7,13 +7,11 @@
 
 import healpy as hp
 import numpy as np
-import numpy.testing as nt
 from astropy import units as u
 
 from .. import ops as ops
 from .. import templates
 from ..accelerator import accel_enabled
-from ..noise import Noise
 from ..observation import default_values as defaults
 from ..pixels import PixelData, PixelDistribution
 from ..pixels_io_healpix import write_healpix_fits
@@ -541,7 +539,6 @@ def test_compare_madam_diagpre(self):
         step_seconds = 5.0
         tmpl = templates.Offset(
             times=defaults.times,
-            det_flags=None,
             noise_model=default_model.noise_model,
             step_time=step_seconds * u.second,
             use_noise_prior=True,
@@ -614,7 +611,6 @@ def test_compare_madam_diagpre(self):
         madam = ops.Madam(
             params=pars,
             det_data=defaults.det_data,
-            det_flags=None,
             pixel_pointing=pixels,
             stokes_weights=weights,
             noise_model="noise_model",
@@ -771,7 +767,6 @@ def test_compare_madam_bandpre(self):
         step_seconds = 5.0
         tmpl = templates.Offset(
             times=defaults.times,
-            det_flags=None,
             noise_model=default_model.noise_model,
             step_time=step_seconds * u.second,
             use_noise_prior=True,
@@ -844,7 +839,6 @@ def test_compare_madam_bandpre(self):
         madam = ops.Madam(
             params=pars,
             det_data=defaults.det_data,
-            det_flags=None,
             pixel_pointing=pixels,
             stokes_weights=weights,
             noise_model="noise_model",
diff --git a/src/toast/tests/ops_mapmaker_solve.py b/src/toast/tests/ops_mapmaker_solve.py
index 4213f45b0..21c642c5c 100644
--- a/src/toast/tests/ops_mapmaker_solve.py
+++ b/src/toast/tests/ops_mapmaker_solve.py
@@ -245,7 +245,6 @@ def test_lhs(self):
         binner.binned = "lhs_binned"
         out_amps = "out_amplitudes"
         data[out_amps] = data["amplitudes"].duplicate()
-
         lhs_calc = SolverLHS(
             det_temp=defaults.det_data,
             binning=binner,
diff --git a/src/toast/tests/ops_pointing_healpix.py b/src/toast/tests/ops_pointing_healpix.py
index 4685ddf43..2ebbc2f0c 100644
--- a/src/toast/tests/ops_pointing_healpix.py
+++ b/src/toast/tests/ops_pointing_healpix.py
@@ -10,6 +10,7 @@
 from .. import ops as ops
 from .. import qarray as qa
 from .._libtoast import healpix_pixels, stokes_weights
+from ..accelerator import ImplementationType
 from ..healpix import HealpixPixels
 from ..intervals import IntervalList, interval_dtype
 from ..observation import default_values as defaults
@@ -219,9 +220,9 @@ def test_hpix_simple(self):
         pixels.apply(data)
 
         # Also make a copy using a python codepath
-        detpointing.use_python = True
+        detpointing.kernel_implementation = ImplementationType.NUMPY
         detpointing.quats = "pyquat"
-        pixels.use_python = True
+        pixels.kernel_implementation = ImplementationType.NUMPY
         pixels.quats = "pyquat"
         pixels.pixels = "pypix"
         pixels.apply(data)
@@ -262,13 +263,13 @@ def test_pixweight_pipe(self):
         pipe.apply(data)
 
         # Also make a copy using a python codepath
-        detpointing.use_python = True
+        detpointing.kernel_implementation = ImplementationType.NUMPY
         detpointing.quats = "pyquat"
-        pixels.use_python = True
+        pixels.kernel_implementation = ImplementationType.NUMPY
         pixels.quats = "pyquat"
         pixels.pixels = "pypixels"
         pixels.apply(data)
-        weights.use_python = True
+        weights.kernel_implementation = ImplementationType.NUMPY
         weights.quats = "pyquat"
         weights.weights = "pyweight"
         weights.apply(data)
diff --git a/src/toast/tests/ops_polyfilter.py b/src/toast/tests/ops_polyfilter.py
index 2ae40aaf4..8d7aecd85 100644
--- a/src/toast/tests/ops_polyfilter.py
+++ b/src/toast/tests/ops_polyfilter.py
@@ -11,6 +11,7 @@
 
 from .. import ops as ops
 from .. import qarray as qa
+from ..accelerator import ImplementationType
 from ..noise import Noise
 from ..observation import default_values as defaults
 from ..pixels import PixelData, PixelDistribution
@@ -358,7 +359,7 @@ def test_polyfilter2D(self):
             shared_flag_mask=255,
             focalplane_key="wafer",
         )
-        polyfilter.use_python = True
+        polyfilter.kernel_implementation = ImplementationType.NUMPY
         polyfilter.apply(data)
 
         # Plot filtered TOD
@@ -380,8 +381,8 @@ def test_polyfilter2D(self):
         # Do the same with C++ implementation
 
         polyfilter.det_data = defaults.det_data
+        polyfilter.kernel_implementation = ImplementationType.DEFAULT
         polyfilter.det_flags = defaults.det_flags
-        polyfilter.use_python = False
         polyfilter.apply(data)
 
         if data.comm.world_rank == 0:
@@ -402,11 +403,14 @@ def test_polyfilter2D(self):
         # Check for consistency
         for ob in data.obs:
             for det in ob.local_detectors:
-                self.assertTrue(
-                    np.allclose(
-                        ob.detdata[defaults.det_data][det], ob.detdata["pyfilter"][det]
-                    )
+                check = np.allclose(
+                    ob.detdata[defaults.det_data][det], ob.detdata["pyfilter"][det]
                 )
+                if not check:
+                    print(
+                        f"{ob.detdata[defaults.det_data][det]} != {ob.detdata['pyfilter'][det]}"
+                    )
+                    self.assertTrue(False)
                 self.assertTrue(
                     np.all(
                         np.equal(
diff --git a/src/toast/tests/ops_scan_map.py b/src/toast/tests/ops_scan_map.py
index f9584a984..6c3f3b143 100644
--- a/src/toast/tests/ops_scan_map.py
+++ b/src/toast/tests/ops_scan_map.py
@@ -9,6 +9,7 @@
 from astropy import units as u
 
 from .. import ops as ops
+from ..accelerator import ImplementationType
 from ..observation import default_values as defaults
 from ..pixels import PixelData
 from ._helpers import close_data, create_fake_sky, create_outdir, create_satellite_data
@@ -42,6 +43,7 @@ def test_scan(self):
 
         # Create fake polarized sky pixel values locally
         create_fake_sky(data, "pixel_dist", "fake_map")
+        map_data = data["fake_map"]
 
         # Scan map into timestreams
         scanner = ops.ScanMap(
@@ -50,26 +52,30 @@ def test_scan(self):
             weights=weights.weights,
             map_key="fake_map",
         )
-        scanner.apply(data)
-
-        # Manual check of the projection of map values to timestream
 
-        map_data = data["fake_map"]
-        for ob in data.obs:
-            for det in ob.local_detectors:
-                wt = ob.detdata[weights.weights][det]
-                local_sm, local_pix = data["pixel_dist"].global_pixel_to_submap(
-                    ob.detdata[pixels.pixels][det]
-                )
-                for i in range(ob.n_local_samples):
-                    if local_pix[i] < 0:
-                        continue
-                    val = 0.0
-                    for j in range(3):
-                        val += wt[i, j] * map_data.data[local_sm[i], local_pix[i], j]
-                    np.testing.assert_almost_equal(
-                        val, ob.detdata[defaults.det_data][det, i]
+        for impl in [ImplementationType.DEFAULT, ImplementationType.NUMPY]:
+            scanner.zero = True
+            scanner.kernel_implementation = impl
+            scanner.apply(data)
+
+            # Manual check of the projection of map values to timestream
+            for ob in data.obs:
+                for det in ob.local_detectors:
+                    wt = ob.detdata[weights.weights][det]
+                    local_sm, local_pix = data["pixel_dist"].global_pixel_to_submap(
+                        ob.detdata[pixels.pixels][det]
                     )
+                    for i in range(ob.n_local_samples):
+                        if local_pix[i] < 0:
+                            continue
+                        val = 0.0
+                        for j in range(3):
+                            val += (
+                                wt[i, j] * map_data.data[local_sm[i], local_pix[i], j]
+                            )
+                        np.testing.assert_almost_equal(
+                            val, ob.detdata[defaults.det_data][det, i]
+                        )
 
         close_data(data)
 
diff --git a/src/toast/tests/template_offset.py b/src/toast/tests/template_offset.py
index e0894e2cf..82d1a9740 100644
--- a/src/toast/tests/template_offset.py
+++ b/src/toast/tests/template_offset.py
@@ -9,6 +9,7 @@
 from astropy import units as u
 
 from .. import ops
+from ..accelerator import ImplementationType
 from ..observation import default_values as defaults
 from ..templates import Offset
 from ..utils import rate_from_times
@@ -119,7 +120,7 @@ def test_compare(self):
             times=defaults.times,
             noise_model=noise_model.noise_model,
             step_time=step_seconds * u.second,
-            use_python=True,
+            kernel_implementation=ImplementationType.NUMPY,
         )
 
         # Set the data