added xp to key functions

Author: Jammy2211

autogalaxy/galaxy/galaxies.py

@@ -119,7 +119,7 @@ def image_2d_from(
         )
 
     def galaxy_image_2d_dict_from(
-        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None,
     ) -> {Galaxy: np.ndarray}:
         """
         Returns a dictionary associating every `Galaxy` object with its corresponding 2D image, using the instance

autogalaxy/imaging/fit_imaging.py

@@ -185,6 +185,7 @@ def galaxy_image_dict(self) -> Dict[Galaxy, np.ndarray]:
 
         galaxy_image_2d_dict = self.galaxies.galaxy_image_2d_dict_from(
             grid=self.grids.lp,
+            xp=self._xp
         )
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
@@ -212,6 +213,7 @@ def galaxy_model_image_dict(self) -> Dict[Galaxy, np.ndarray]:
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
+            xp=self._xp
         )
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(

autogalaxy/interferometer/fit_interferometer.py

@@ -161,7 +161,7 @@ def galaxy_image_dict(self) -> Dict[Galaxy, np.ndarray]:
         For modeling, this dictionary is used to set up the `adapt_images` that adapt certain pixelizations to the
         data being fitted.
         """
-        galaxy_image_dict = self.galaxies.galaxy_image_2d_dict_from(grid=self.grids.lp)
+        galaxy_image_dict = self.galaxies.galaxy_image_2d_dict_from(grid=self.grids.lp, xp=self._xp)
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
             use_operated=False
@@ -184,7 +184,7 @@ def galaxy_model_visibilities_dict(self) -> Dict[Galaxy, np.ndarray]:
         data being fitted.
         """
         galaxy_model_visibilities_dict = self.galaxies.galaxy_visibilities_dict_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
         )
 
         galaxy_linear_obj_data_dict = self.galaxy_linear_obj_data_dict_from(

autogalaxy/operate/image.py

@@ -363,12 +363,12 @@ class OperateImageGalaxies(OperateImageList):
     """
 
     def galaxy_image_2d_dict_from(
-        self, grid: aa.Grid2D, operated_only: Optional[bool] = None
+        self, grid: aa.Grid2D, xp=np, operated_only: Optional[bool] = None
     ) -> Dict[Galaxy, aa.Array2D]:
         raise NotImplementedError
 
     def galaxy_blurred_image_2d_dict_from(
-        self, grid, psf, blurring_grid
+        self, grid, psf, blurring_grid, xp=np
     ) -> Dict[Galaxy, aa.Array2D]:
         """
         Evaluate the light object's dictionary mapping galaixes to their corresponding 2D images and convolve each
@@ -392,15 +392,15 @@ def galaxy_blurred_image_2d_dict_from(
         """
 
         galaxy_image_2d_not_operated_dict = self.galaxy_image_2d_dict_from(
-            grid=grid, operated_only=False
+            grid=grid, operated_only=False, xp=xp
         )
 
         galaxy_blurring_image_2d_not_operated_dict = self.galaxy_image_2d_dict_from(
-            grid=blurring_grid, operated_only=False
+            grid=blurring_grid, operated_only=False, xp=xp
         )
 
         galaxy_image_2d_operated_dict = self.galaxy_image_2d_dict_from(
-            grid=grid, operated_only=True
+            grid=grid, operated_only=True, xp=xp
         )
 
         galaxy_blurred_image_2d_dict = {}
@@ -414,6 +414,7 @@ def galaxy_blurred_image_2d_dict_from(
             blurred_image_2d = psf.convolved_image_from(
                 image=image_2d_not_operated,
                 blurring_image=blurring_image_2d_not_operated,
+                xp=xp
             )
 
             image_2d_operated = galaxy_image_2d_operated_dict[galaxy_key]
@@ -451,7 +452,7 @@ def galaxy_visibilities_dict_from(
             in the uv-plane.
         """
 
-        galaxy_image_2d_dict = self.galaxy_image_2d_dict_from(grid=grid)
+        galaxy_image_2d_dict = self.galaxy_image_2d_dict_from(grid=grid, xp=xp)
 
         galaxy_visibilities_dict = {}