Enable jit on pad and crop operators

lanl · Feb 25, 2024 · 203b3c9 · 203b3c9
1 parent 2e92b2e
commit 203b3c9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scico/functional/_tvnorm.py b/scico/functional/_tvnorm.py
@@ -137,9 +137,9 @@ def prox(self, v: Array, lam: float = 1.0, **kwargs) -> Array:
         if self.W is None or self.W.shape[1] != w_input_shape:
             self.W = self._haar_operator(ndims, w_input_shape, v.dtype)
             if not self.circular:
-                P = Pad(v.shape, pad_width=(((0, 1),) * ndims), mode="edge")
+                P = Pad(v.shape, pad_width=(((0, 1),) * ndims), mode="edge", jit=True)
                 self.WP = self.W @ P
-                C = Crop(crop_width=(((0, 1),) * ndims), input_shape=w_input_shape)
+                C = Crop(crop_width=(((0, 1),) * ndims), input_shape=w_input_shape, jit=True)
                 self.CWT = C @ self.W.T
 
         if self.circular: