diff --git a/NDSL b/NDSL
index 2368e5d8..3d8fa05e 160000
--- a/NDSL
+++ b/NDSL
@@ -1 +1 @@
-Subproject commit 2368e5d8bc2cb3ddeb4395f33f7f33c35a59e18b
+Subproject commit 3d8fa05ee7b78182afd2e749532582ffaa25b595
diff --git a/pace/diagnostics.py b/pace/diagnostics.py
index 35895e03..6afdc768 100644
--- a/pace/diagnostics.py
+++ b/pace/diagnostics.py
@@ -58,6 +58,7 @@ def select_data(self, state: DycoreState):
                 origin=getattr(state, name).origin[0:2],
                 extent=getattr(state, name).extent[0:2],
                 units=getattr(state, name).units,
+                backend=getattr(state, name).backend,
             )
         return output
 
@@ -235,7 +236,7 @@ def cleanup(self):
         pass
 
 
-def _compute_column_integral(name: str, q_in: Quantity, delp: Quantity):
+def _compute_column_integral(name: str, q_in: Quantity, delp: Quantity) -> Quantity:
     """
     Compute column integrated mixing ratio (e.g., total liquid water path)
 
@@ -243,19 +244,25 @@ def _compute_column_integral(name: str, q_in: Quantity, delp: Quantity):
         name: name of the tracer
         q_in: tracer mixing ratio
         delp: pressure thickness of atmospheric layer
+
+    Returns:
+        The column integral.
     """
-    assert len(q_in.dims) > 2
+    if len(q_in.dims) <= 2:
+        raise RuntimeError("This function assumes that q_in is at least 3-dimensional.")
+
     if q_in.dims[2] != Z_DIM:
-        raise NotImplementedError(
-            "this function assumes the z-dimension is the third dimension"
+        raise RuntimeError(
+            "This function assumes the z-dimension is the third dimension"
         )
+
     k_slice = slice(q_in.origin[2], q_in.origin[2] + q_in.extent[2])
-    column_integral = Quantity(
+    return Quantity(
         RGRAV
         * q_in.np.sum(q_in.data[:, :, k_slice] * delp.data[:, :, k_slice], axis=2),
         dims=tuple(q_in.dims[:2]) + tuple(q_in.dims[3:]),
         origin=q_in.metadata.origin[0:2],
         extent=(q_in.metadata.extent[0], q_in.metadata.extent[1]),
         units="kg/m**2",
+        backend=q_in.backend,
     )
-    return column_integral