fix: make jax imports lazy in Grid2DContour properties

contour_list and contour_array did unconditional `import jax.numpy as jnp`
inside their property bodies, which broke the numpy code path on Python
3.9 / 3.10 under the new optional [jax] extra. contour_list only used jax
for an isinstance check (replaced with duck-typing on .array); contour_array's
JAX-array construction is replaced with the numpy equivalent (the downstream
consumer materialises to numpy anyway).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people

autoarray/operators/contour.py

@@ -38,16 +38,14 @@ def contour_array(self):
         if self._contour_array is not None:
             return self._contour_array
 
-        import jax.numpy as jnp
-
         pixel_centres = geometry_util.grid_pixel_centres_2d_slim_from(
             grid_scaled_2d_slim=np.array(self.grid),
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
         ).astype("int")
 
-        arr = jnp.zeros(self.shape_native)
-        arr = arr.at[tuple(jnp.array(pixel_centres).T)].set(1)
+        arr = np.zeros(self.shape_native)
+        arr[tuple(pixel_centres.T)] = 1
 
         return arr
 
@@ -56,12 +54,12 @@ def contour_list(self):
         # make sure to use base numpy to convert JAX array back to a normal array
 
         from skimage import measure
-        import jax.numpy as jnp
 
-        if isinstance(self.contour_array, jnp.ndarray):
-            contour_array = np.array(self.contour_array)
+        contour_array = self.contour_array
+        if hasattr(contour_array, "array"):
+            contour_array = np.array(contour_array.array)
         else:
-            contour_array = np.array(self.contour_array.array)
+            contour_array = np.array(contour_array)
 
         contour_indices_list = measure.find_contours(contour_array, 0)