feat: thread xp through Grid2DIrregular.grid_2d_via_deflection_grid_from

Removes the known JAX-jit blocker for the PyAutoLens point-source
source-plane fit path (PyAutoLabs/PyAutoLens#657, phase 2).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: 2 people

autoarray/structures/grids/irregular_2d.py

@@ -168,7 +168,7 @@ def extent_with_buffer_from(self, buffer: float = 1.0e-8) -> List[float]:
         ]
 
     def grid_2d_via_deflection_grid_from(
-        self, deflection_grid: np.ndarray
+        self, deflection_grid: np.ndarray, xp=np
     ) -> "Grid2DIrregular":
         """
         Returns a new Grid2DIrregular from this grid coordinates, where the (y,x) coordinates of this grid have a
@@ -181,8 +181,13 @@ def grid_2d_via_deflection_grid_from(
         ----------
         deflection_grid
             The grid of (y,x) coordinates which is subtracted from this grid.
+        xp
+            The array module (``numpy`` or ``jax.numpy``) used to construct the returned grid, mirroring
+            ``subtracted_from`` / ``subtracted_and_rotated_from``. Passed through explicitly by the caller rather
+            than inferred from ``self._xp``, so JIT-traced call sites (where ``self`` may not carry a reliable
+            ``use_jax`` flag) do not silently fall back to NumPy.
         """
-        return Grid2DIrregular(values=self - deflection_grid, xp=self._xp)
+        return Grid2DIrregular(values=self.array - xp.asarray(deflection_grid), xp=xp)
 
     def squared_distances_to_coordinate_from(
         self, coordinate: Tuple[float, float] = (0.0, 0.0)