diff --git a/autolens/__init__.py b/autolens/__init__.py
index 19d145338..2597bc519 100644
--- a/autolens/__init__.py
+++ b/autolens/__init__.py
@@ -29,6 +29,7 @@
 from autoarray.operators.transformer import TransformerDFT
 from autoarray.operators.transformer import TransformerNUFFT
 from autoarray.preloads import Preloads
+from autoarray.preloads import mapper_indices_from
 from autoarray.structures.arrays.uniform_1d import Array1D
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.arrays.rgb import Array2DRGB
@@ -115,8 +116,6 @@
 from .quantity.fit_quantity import FitQuantity
 from .quantity.model.analysis import AnalysisQuantity
 
-from .analysis.preloads import mapper_indices_from
-
 from . import exc
 from . import mock as m
 from . import util
diff --git a/autolens/aggregator/subplot.py b/autolens/aggregator/subplot.py
index aea258ea0..5ff2a926c 100644
--- a/autolens/aggregator/subplot.py
+++ b/autolens/aggregator/subplot.py
@@ -69,8 +69,8 @@ class SubplotFitX1Plane(Enum):
     signal_to_noise_map = (1, 0)
     model_data = (2, 0)
     lens_light_subtracted_image = (0, 1)
-    lens_light_subtracted_image_zero = (0, 1)
-    normalized_residual_map = (0, 2)
+    lens_light_subtracted_image_zero = (1, 1)
+    normalized_residual_map = (2, 1)
 
 
 class SubplotFit(Enum):
diff --git a/autolens/analysis/analysis/lens.py b/autolens/analysis/analysis/lens.py
index a539cb6fe..df7cd19fa 100644
--- a/autolens/analysis/analysis/lens.py
+++ b/autolens/analysis/analysis/lens.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import logging
 import numpy as np
 from typing import List, Optional
@@ -96,7 +95,7 @@ def tracer_via_instance_from(
         )
 
     def log_likelihood_penalty_from(
-        self, instance: af.ModelInstance
+        self, instance: af.ModelInstance, xp=np
     ) -> Optional[float]:
         """
         Call the positions overwrite log likelihood function, which add a penalty term to the likelihood if the
@@ -117,7 +116,7 @@ def log_likelihood_penalty_from(
         The penalty value of the positions log likelihood, if the positions do not trace close in the source plane,
         else a None is returned to indicate there is no penalty.
         """
-        log_likelihood_penalty = jnp.array(0.0)
+        log_likelihood_penalty = xp.array(0.0)
 
         if self.positions_likelihood_list is not None:
 
@@ -125,7 +124,7 @@ def log_likelihood_penalty_from(
                 for positions_likelihood in self.positions_likelihood_list:
                     log_likelihood_penalty = (
                         positions_likelihood.log_likelihood_penalty_from(
-                            instance=instance, analysis=self
+                            instance=instance, analysis=self, xp=xp
                         )
                     )
 
diff --git a/autolens/analysis/model_util.py b/autolens/analysis/model_util.py
index ca87ecc4b..18566478c 100644
--- a/autolens/analysis/model_util.py
+++ b/autolens/analysis/model_util.py
@@ -1,130 +1,7 @@
-import numpy as np
-from typing import Optional, Tuple
-
 import autofit as af
 import autolens as al
 
-
-def mge_model_from(
-    mask_radius: float,
-    total_gaussians: int = 30,
-    gaussian_per_basis: int = 1,
-    centre_prior_is_uniform: bool = True,
-    centre: Tuple[float, float] = (0.0, 0.0),
-    centre_fixed: Optional[Tuple[float, float]] = None,
-    use_spherical: bool = False,
-) -> af.Collection:
-    """
-    Construct a Multi-Gaussian Expansion (MGE) for the lens or source galaxy light
-
-    This model is designed as a "start here" configuration for lens modeling:
-
-    - The lens and source light are represented by a Basis object composed of many
-      Gaussian light profiles with fixed logarithmically spaced widths (`sigma`).
-    - All Gaussians within each basis share common centres and ellipticity
-      components, reducing degeneracy while retaining flexibility.
-
-    - Users can combine with a lens mass model of their choiuce.
-
-    The resulting model provides a good balance of speed, flexibility, and accuracy
-    for fitting most galaxy-scale strong lenses.
-
-    This code is mostly to make the API simple for new users, hiding the technical
-    details of setting up an MGE. More advanced users may wish to customize the
-    model further.
-
-    Parameters
-    ----------
-    mask_radius
-        The outer radius (in arcseconds) of the circular mask applied to the data.
-        This determines the maximum Gaussian width (`sigma`) used in the lens MGE.
-    lens_total_gaussians
-        Total number of Gaussian light profiles used in the lens MGE basis.
-    source_total_gaussians
-        Total number of Gaussian light profiles used in the source MGE basis.
-    lens_gaussian_per_basis
-        Number of separate Gaussian bases to include for the lens light profile.
-        Each basis has `lens_total_gaussians` components.
-    source_gaussian_per_basis
-        Number of separate Gaussian bases to include for the source light profile.
-        Each basis has `source_total_gaussians` components.
-
-    Returns
-    -------
-    model : af.Collection
-        An `autofit.Collection` containing:
-        - A lens galaxy at redshift 0.5, with:
-          * bulge light profile: MGE basis of Gaussians
-          * mass profile: Isothermal ellipsoid
-          * external shear
-        - A source galaxy at redshift 1.0, with:
-          * bulge light profile: MGE basis of Gaussians
-
-    Notes
-    -----
-    - Lens light Gaussians have widths (sigma) logarithmically spaced between 0.01"
-      and the mask radius.
-    - Source light Gaussians have widths logarithmically spaced between 0.01" and 1.0".
-    - Gaussian centres are free parameters but tied across all components in each
-      basis to reduce dimensionality.
-    - This function is a convenience utility: it hides the technical setup of MGE
-      composition and provides a ready-to-use lens model for quick experimentation.
-    """
-
-    # The sigma values of the Gaussians will be fixed to values spanning 0.01 to the mask radius, 3.0".
-    log10_sigma_list = np.linspace(-4, np.log10(mask_radius), total_gaussians)
-
-    # By defining the centre here, it creates two free parameters that are assigned below to all Gaussians.
-
-    if centre_fixed is not None:
-        centre_0 = centre[0]
-        centre_1 = centre[1]
-    elif centre_prior_is_uniform:
-        centre_0 = af.UniformPrior(
-            lower_limit=centre[0] - 0.1, upper_limit=centre[0] + 0.1
-        )
-        centre_1 = af.UniformPrior(
-            lower_limit=centre[1] - 0.1, upper_limit=centre[1] + 0.1
-        )
-    else:
-        centre_0 = af.GaussianPrior(mean=centre[0], sigma=0.3)
-        centre_1 = af.GaussianPrior(mean=centre[1], sigma=0.3)
-
-    if use_spherical:
-        model_cls = al.lp_linear.GaussianSph
-    else:
-        model_cls = al.lp_linear.Gaussian
-
-    bulge_gaussian_list = []
-
-    for j in range(gaussian_per_basis):
-        # A list of Gaussian model components whose parameters are customized belows.
-
-        gaussian_list = af.Collection(
-            af.Model(model_cls) for _ in range(total_gaussians)
-        )
-
-        # Iterate over every Gaussian and customize its parameters.
-
-        for i, gaussian in enumerate(gaussian_list):
-            gaussian.centre.centre_0 = centre_0  # All Gaussians have same y centre.
-            gaussian.centre.centre_1 = centre_1  # All Gaussians have same x centre.
-            if not use_spherical:
-                gaussian.ell_comps = gaussian_list[
-                    0
-                ].ell_comps  # All Gaussians have same elliptical components.
-            gaussian.sigma = (
-                10 ** log10_sigma_list[i]
-            )  # All Gaussian sigmas are fixed to values above.
-
-        bulge_gaussian_list += gaussian_list
-
-    # The Basis object groups many light profiles together into a single model component.
-
-    return af.Model(
-        al.lp_basis.Basis,
-        profile_list=bulge_gaussian_list,
-    )
+from autogalaxy.analysis.model_util import mge_model_from
 
 
 def simulator_start_here_model_from(
diff --git a/autolens/analysis/positions.py b/autolens/analysis/positions.py
index 79683e266..8d7d96d7e 100644
--- a/autolens/analysis/positions.py
+++ b/autolens/analysis/positions.py
@@ -1,10 +1,7 @@
 import jax
-import jax.numpy as jnp
 import numpy as np
-from typing import Optional, Union
+from typing import Optional
 from os import path
-import os
-from typing import Dict
 
 import autoarray as aa
 import autofit as af
@@ -137,8 +134,8 @@ def output_positions_info(
             f.write("")
 
     def log_likelihood_penalty_from(
-        self, instance: af.ModelInstance, analysis: AnalysisDataset
-    ) -> jnp.array:
+        self, instance: af.ModelInstance, analysis: AnalysisDataset, xp=np
+    ) -> np.array:
         """
         Returns a log-likelihood penalty used to constrain lens models where multiple image-plane
         positions do not trace to within a threshold distance of one another in the source-plane.
@@ -176,16 +173,17 @@ def log_likelihood_penalty_from(
         tracer = analysis.tracer_via_instance_from(instance=instance)
 
         if not tracer.has(cls=ag.mp.MassProfile) or len(tracer.planes) == 1:
-            return (jnp.array(0.0),)
+            return xp.array(0.0)
 
         positions_fit = SourceMaxSeparation(
             data=self.positions,
             noise_map=None,
             tracer=tracer,
             plane_redshift=self.plane_redshift,
+            xp=xp
         )
 
-        max_separation = jnp.max(
+        max_separation = xp.max(
             positions_fit.furthest_separations_of_plane_positions.array
         )
 
@@ -194,5 +192,5 @@ def log_likelihood_penalty_from(
         return jax.lax.cond(
             max_separation > self.threshold,
             lambda: penalty,
-            lambda: jnp.array(0.0),
+            lambda: xp.array(0.0),
         )
diff --git a/autolens/analysis/preloads.py b/autolens/analysis/preloads.py
deleted file mode 100644
index fb24761a9..000000000
--- a/autolens/analysis/preloads.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import jax.numpy as jnp
-
-
-def mapper_indices_from(total_linear_light_profiles, total_mapper_pixels):
-    return jnp.arange(
-        total_linear_light_profiles,
-        total_linear_light_profiles + total_mapper_pixels,
-        dtype=int,
-    )
diff --git a/autolens/config/visualize/plots.yaml b/autolens/config/visualize/plots.yaml
index 84608a529..4adb8c692 100644
--- a/autolens/config/visualize/plots.yaml
+++ b/autolens/config/visualize/plots.yaml
@@ -66,5 +66,3 @@ fit_quantity:                              # Settings for plots of fit quantitie
 galaxies:                                  # Settings for plots of galaxies (e.g. GalaxiesPlotter).
   subplot_galaxies: true                   # Plot subplot of all quantities in each galaxies group (e.g. images, convergence)?
   subplot_galaxy_images: false             # Plot subplot of the image of each galaxy in the model?
-  subplot_galaxies_1d: false               # Plot subplot of all quantities in 1D of each galaxies group (e.g. images, convergence)?
-  subplot_galaxies_1d_decomposed: false    # Plot subplot of all quantities in 1D decomposed of each galaxies group (e.g. images, convergence)?
diff --git a/autolens/exc.py b/autolens/exc.py
index 83a73449a..5f49ff037 100644
--- a/autolens/exc.py
+++ b/autolens/exc.py
@@ -32,7 +32,7 @@ class PixelizationException(af.exc.FitException):
     """
     Raises exceptions associated with the `inversion/pixelization` modules and `Pixelization` classes.
 
-    For example if a `Rectangular` mesh has dimensions below 3x3.
+    For example if a `RectangularMagnification` mesh has dimensions below 3x3.
 
     This exception overwrites `autoarray.exc.PixelizationException` in order to add a `FitException`. This means that
     if this exception is raised during a model-fit in the analysis class's `log_likelihood_function` that model
diff --git a/autolens/imaging/fit_imaging.py b/autolens/imaging/fit_imaging.py
index bd55a9a5a..7eb9bd4e1 100644
--- a/autolens/imaging/fit_imaging.py
+++ b/autolens/imaging/fit_imaging.py
@@ -24,6 +24,7 @@ def __init__(
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = None,
+        xp=np
     ):
         """
         Fits an imaging dataset using a `Tracer` object.
@@ -64,7 +65,7 @@ def __init__(
             Settings controlling how an inversion is fitted for example which linear algebra formalism is used.
         """
 
-        super().__init__(dataset=dataset, dataset_model=dataset_model)
+        super().__init__(dataset=dataset, dataset_model=dataset_model, xp=xp)
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
@@ -86,12 +87,14 @@ def blurred_image(self) -> aa.Array2D:
         ):
             return self.tracer.image_2d_from(
                 grid=self.grids.lp,
+                xp=self._xp,
             )
 
         return self.tracer.blurred_image_2d_from(
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
+            xp=self._xp,
         )
 
     @property
@@ -118,6 +121,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=self._xp,
         )
 
     @cached_property
diff --git a/autolens/imaging/model/analysis.py b/autolens/imaging/model/analysis.py
index fe4d3a1a7..a73125e54 100644
--- a/autolens/imaging/model/analysis.py
+++ b/autolens/imaging/model/analysis.py
@@ -1,5 +1,5 @@
 import logging
-
+import numpy as np
 
 import autofit as af
 import autogalaxy as ag
@@ -46,7 +46,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance: af.ModelInstance) -> float:
+    def log_likelihood_function(self, instance: af.ModelInstance, xp=np) -> float:
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the imaging dataset.
@@ -86,14 +86,15 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
         """
 
         log_likelihood_penalty = self.log_likelihood_penalty_from(
-            instance=instance
+            instance=instance,
+            xp=xp
         )
 
-        return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
+        return self.fit_from(instance=instance, xp=xp).figure_of_merit - log_likelihood_penalty
 
     def fit_from(
         self,
-        instance: af.ModelInstance,
+        instance: af.ModelInstance, xp=np
     ) -> FitImaging:
         """
         Given a model instance create a `FitImaging` object.
@@ -130,7 +131,8 @@ def fit_from(
             dataset_model=dataset_model,
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
-            preloads=self.preloads
+            preloads=self.preloads,
+            xp=xp
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autolens/interferometer/fit_interferometer.py b/autolens/interferometer/fit_interferometer.py
index e4993d891..6cae142b2 100644
--- a/autolens/interferometer/fit_interferometer.py
+++ b/autolens/interferometer/fit_interferometer.py
@@ -21,6 +21,7 @@ def __init__(
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = None,
+        xp=np,
     ):
         """
         Fits an interferometer dataset using a `Tracer` object.
@@ -76,12 +77,14 @@ def __init__(
         super().__init__(
             dataset=dataset,
             dataset_model=dataset_model,
+            xp=xp,
         )
         AbstractFitInversion.__init__(
             self=self, model_obj=tracer, settings_inversion=settings_inversion
         )
 
         self.preloads = preloads
+        self._xp = xp
 
     @property
     def profile_visibilities(self) -> aa.Visibilities:
@@ -90,7 +93,7 @@ def profile_visibilities(self) -> aa.Visibilities:
         transform to the sum of light profile images.
         """
         return self.tracer.visibilities_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
         )
 
     @property
@@ -117,6 +120,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             adapt_images=self.adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=self._xp,
         )
 
     @cached_property
diff --git a/autolens/interferometer/model/analysis.py b/autolens/interferometer/model/analysis.py
index 9257aedd6..947214158 100644
--- a/autolens/interferometer/model/analysis.py
+++ b/autolens/interferometer/model/analysis.py
@@ -114,7 +114,7 @@ def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
 
         return self
 
-    def log_likelihood_function(self, instance):
+    def log_likelihood_function(self, instance, xp=np):
         """
         Given an instance of the model, where the model parameters are set via a non-linear search, fit the model
         instance to the interferometer dataset.
@@ -153,14 +153,16 @@ def log_likelihood_function(self, instance):
             The log likelihood indicating how well this model instance fitted the interferometer data.
         """
 
-        log_likelihood_penalty = self.log_likelihood_penalty_from(instance=instance)
+        log_likelihood_penalty = self.log_likelihood_penalty_from(
+            instance=instance, xp=xp
+        )
 
-        return self.fit_from(instance=instance).figure_of_merit - log_likelihood_penalty
+        return (
+            self.fit_from(instance=instance, xp=xp).figure_of_merit
+            - log_likelihood_penalty
+        )
 
-    def fit_from(
-        self,
-        instance: af.ModelInstance,
-    ) -> FitInterferometer:
+    def fit_from(self, instance: af.ModelInstance, xp=np) -> FitInterferometer:
         """
         Given a model instance create a `FitInterferometer` object.
 
@@ -196,6 +198,7 @@ def fit_from(
             adapt_images=adapt_images,
             settings_inversion=self.settings_inversion,
             preloads=self.preloads,
+            xp=xp,
         )
 
     def save_attributes(self, paths: af.DirectoryPaths):
diff --git a/autolens/lens/mock/mock_tracer.py b/autolens/lens/mock/mock_tracer.py
index f97349712..d6ca0c340 100644
--- a/autolens/lens/mock/mock_tracer.py
+++ b/autolens/lens/mock/mock_tracer.py
@@ -1,3 +1,6 @@
+import numpy as np
+
+
 class MockTracer:
     def __init__(
         self, traced_grid_2d_list_from=None, image_plane_mesh_grid_pg_list=None
@@ -5,7 +8,7 @@ def __init__(
         self.image_plane_mesh_grid_pg_list = image_plane_mesh_grid_pg_list
         self._traced_grid_2d_list_from = traced_grid_2d_list_from
 
-    def traced_grid_2d_list_from(self, grid):
+    def traced_grid_2d_list_from(self, grid, xp=np):
         return self._traced_grid_2d_list_from
 
     def plane_index_via_redshift_from(self, redshift):
@@ -61,5 +64,5 @@ def einstein_radius_from(self, grid):
     def einstein_mass_angular_from(self, grid):
         return self.einstein_mass
 
-    def time_delays_from(self, grid):
+    def time_delays_from(self, grid, xp=np):
         return self.time_delays
diff --git a/autolens/lens/to_inversion.py b/autolens/lens/to_inversion.py
index 4b7a27ec7..46c1430db 100644
--- a/autolens/lens/to_inversion.py
+++ b/autolens/lens/to_inversion.py
@@ -18,6 +18,7 @@ def __init__(
         adapt_images: Optional[ag.AdaptImages] = None,
         settings_inversion: aa.SettingsInversion = aa.SettingsInversion(),
         preloads: aa.Preloads = None,
+        xp=np,
     ):
         """
         Interfaces a dataset and tracer with the inversion module, to setup a linear algebra calculation.
@@ -60,6 +61,7 @@ def __init__(
             dataset=dataset,
             adapt_images=adapt_images,
             settings_inversion=settings_inversion,
+            xp=xp,
         )
 
     @property
@@ -113,7 +115,7 @@ def traced_grid_2d_list_of_inversion(self) -> List[aa.type.Grid2DLike]:
         The traced grids of the inversion, which are cached for efficiency.
         """
         return self.tracer.traced_grid_2d_list_from(
-            grid=self.dataset.grids.pixelization
+            grid=self.dataset.grids.pixelization, xp=self._xp
         )
 
     @cached_property
@@ -155,12 +157,12 @@ def lp_linear_func_list_galaxy_dict(
         lp_linear_galaxy_dict_list = {}
 
         traced_grids_of_planes_list = self.tracer.traced_grid_2d_list_from(
-            grid=self.dataset.grids.lp
+            grid=self.dataset.grids.lp, xp=self._xp
         )
 
         if self.dataset.grids.blurring is not None:
             traced_blurring_grids_of_planes_list = self.tracer.traced_grid_2d_list_from(
-                grid=self.dataset.grids.blurring
+                grid=self.dataset.grids.blurring, xp=self._xp
             )
         else:
             traced_blurring_grids_of_planes_list = [None] * len(
@@ -192,6 +194,7 @@ def lp_linear_func_list_galaxy_dict(
                 galaxies=galaxies,
                 settings_inversion=self.settings_inversion,
                 adapt_images=self.adapt_images,
+                xp=self._xp,
             )
 
             lp_linear_galaxy_dict_of_plane = (
@@ -410,6 +413,7 @@ def mapper_galaxy_dict(self) -> Dict[aa.AbstractMapper, ag.Galaxy]:
                     galaxies=galaxies,
                     adapt_images=self.adapt_images,
                     settings_inversion=self.settings_inversion,
+                    xp=self._xp,
                 )
 
                 galaxies_with_pixelization_list = galaxies.galaxies_with_cls_list_from(
@@ -474,6 +478,7 @@ def inversion(self):
             linear_obj_list=self.linear_obj_list,
             settings=self.settings_inversion,
             preloads=self.preloads,
+            xp=self._xp,
         )
 
         inversion.linear_obj_galaxy_dict = self.linear_obj_galaxy_dict
diff --git a/autolens/lens/tracer.py b/autolens/lens/tracer.py
index d2f621baa..8f025c908 100644
--- a/autolens/lens/tracer.py
+++ b/autolens/lens/tracer.py
@@ -212,7 +212,7 @@ def total_planes(self) -> int:
 
     @aa.grid_dec.to_grid
     def traced_grid_2d_list_from(
-        self, grid: aa.type.Grid2DLike, plane_index_limit: int = Optional[None]
+        self, grid: aa.type.Grid2DLike, xp=np, plane_index_limit: int = Optional[None]
     ) -> List[aa.type.Grid2DLike]:
         """
         Returns a ray-traced grid of 2D Cartesian (y,x) coordinates which accounts for multi-plane ray-tracing.
@@ -262,6 +262,7 @@ def traced_grid_2d_list_from(
             grid=grid,
             cosmology=self.cosmology,
             plane_index_limit=plane_index_limit,
+            xp=xp,
         )
 
         if isinstance(grid, aa.Grid2D):
@@ -270,6 +271,7 @@ def traced_grid_2d_list_from(
                 grid=grid.over_sampled,
                 cosmology=self.cosmology,
                 plane_index_limit=plane_index_limit,
+                xp=xp,
             )
 
             grid_2d_new_list = []
@@ -365,6 +367,7 @@ def upper_plane_index_with_light_profile(self) -> int:
     def image_2d_list_from(
         self,
         grid: aa.type.Grid2DLike,
+        xp=np,
         operated_only: Optional[bool] = None,
     ) -> List[aa.Array2D]:
         """
@@ -410,7 +413,9 @@ def image_2d_list_from(
         """
 
         traced_grid_list = self.traced_grid_2d_list_from(
-            grid=grid, plane_index_limit=self.upper_plane_index_with_light_profile
+            grid=grid,
+            xp=xp,
+            plane_index_limit=self.upper_plane_index_with_light_profile,
         )
 
         image_2d_list = []
@@ -423,6 +428,7 @@ def image_2d_list_from(
                     galaxy.image_2d_from(
                         grid=traced_grid_list[plane_index],
                         operated_only=operated_only,
+                        xp=xp,
                     )
                     for galaxy in galaxies
                 ]
@@ -449,6 +455,7 @@ def image_2d_list_from(
     def image_2d_from(
         self,
         grid: aa.type.Grid2DLike,
+        xp=np,
         operated_only: Optional[bool] = None,
     ) -> aa.Array2D:
         """
@@ -470,12 +477,15 @@ def image_2d_from(
             apply these operations to the images, which may have the `operated_only` input passed to them. This input
             therefore is used to pass the `operated_only` input to these methods.
         """
-        return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
+        return sum(
+            self.image_2d_list_from(grid=grid, operated_only=operated_only, xp=xp)
+        )
 
     def image_2d_via_input_plane_image_from(
         self,
         grid: aa.type.Grid2DLike,
         plane_image: aa.Array2D,
+        xp=np,
         plane_index: int = -1,
         include_other_planes: bool = True,
     ) -> aa.Array2D:
@@ -538,7 +548,7 @@ def image_2d_via_input_plane_image_from(
         )
 
         traced_grid = self.traced_grid_2d_list_from(
-            grid=grid, plane_index_limit=plane_index
+            grid=grid, plane_index_limit=plane_index, xp=xp
         )[plane_index]
 
         image = griddata(
@@ -550,10 +560,10 @@ def image_2d_via_input_plane_image_from(
         )
 
         if isinstance(grid, aa.Grid2D):
-            image = grid.over_sampler.binned_array_2d_from(array=image)
+            image = grid.over_sampler.binned_array_2d_from(array=image, xp=xp)
 
         if include_other_planes:
-            image_list = self.image_2d_list_from(grid=grid, operated_only=False)
+            image_list = self.image_2d_list_from(grid=grid, xp=xp, operated_only=False)
 
             if plane_index < 0:
                 plane_index = self.total_planes + plane_index
@@ -562,10 +572,7 @@ def image_2d_via_input_plane_image_from(
                 if plane_lp_index != plane_index:
                     image += image_list[plane_lp_index]
 
-        return aa.Array2D(
-            values=image,
-            mask=grid.mask,
-        )
+        return aa.Array2D(values=image, mask=grid.mask, xp=xp)
 
     def galaxy_image_2d_dict_from(
         self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
@@ -611,7 +618,7 @@ def galaxy_image_2d_dict_from(
 
     @aa.grid_dec.to_vector_yx
     def deflections_yx_2d_from(
-        self, grid: aa.type.Grid2DLike
+        self, grid: aa.type.Grid2DLike, xp=np
     ) -> Union[aa.VectorYX2D, aa.VectorYX2DIrregular]:
         """
         Returns the 2D deflection angles of all galaxies in the tracer, from the image-plane to the source-plane,
@@ -637,12 +644,12 @@ def deflections_yx_2d_from(
             The 2D (y, x) coordinates where values of the deflections are evaluated.
         """
         if self.total_planes > 1:
-            return self.deflections_between_planes_from(grid=grid)
-        return self.deflections_of_planes_summed_from(grid=grid)
+            return self.deflections_between_planes_from(grid=grid, xp=xp)
+        return self.deflections_of_planes_summed_from(grid=grid, xp=xp)
 
     @aa.grid_dec.to_vector_yx
     def deflections_of_planes_summed_from(
-        self, grid: aa.type.Grid2DLike
+        self, grid: aa.type.Grid2DLike, xp=np
     ) -> Union[aa.VectorYX2D, aa.VectorYX2DIrregular]:
         """
         Returns the summed 2D deflections angles of all galaxies in the tracer, not accounting for multi-plane ray
@@ -671,12 +678,15 @@ def deflections_of_planes_summed_from(
             The 2D (y, x) coordinates where values of the deflections are evaluated.
         """
         return sum(
-            [galaxy.deflections_yx_2d_from(grid=grid) for galaxy in self.galaxies]
+            [
+                galaxy.deflections_yx_2d_from(grid=grid, xp=xp)
+                for galaxy in self.galaxies
+            ]
         )
 
     @aa.grid_dec.to_vector_yx
     def deflections_between_planes_from(
-        self, grid: aa.type.Grid2DLike, plane_i=0, plane_j=-1
+        self, grid: aa.type.Grid2DLike, xp=np, plane_i=0, plane_j=-1
     ) -> Union[aa.VectorYX2D, aa.VectorYX2DIrregular]:
         """
         Returns the summed 2D deflections angles between two input planes in the tracer, accounting for multi-plane
@@ -704,7 +714,7 @@ def deflections_between_planes_from(
         return traced_grids_list[plane_i] - traced_grids_list[plane_j]
 
     @aa.grid_dec.to_array
-    def convergence_2d_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
+    def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D:
         """
         Returns the summed 2D convergence of all galaxies in the tracer from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -726,10 +736,12 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
         grid
             The 2D (y, x) coordinates where values of the convergence are evaluated.
         """
-        return sum([galaxy.convergence_2d_from(grid=grid) for galaxy in self.galaxies])
+        return sum(
+            [galaxy.convergence_2d_from(grid=grid, xp=xp) for galaxy in self.galaxies]
+        )
 
     @aa.grid_dec.to_array
-    def potential_2d_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
+    def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D:
         """
         Returns the summed 2D potential of all galaxies in the tracer from a 2D grid of Cartesian (y,x) coordinates.
 
@@ -751,10 +763,12 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
         grid
             The 2D (y, x) coordinates where values of the potential are evaluated.
         """
-        return sum([galaxy.potential_2d_from(grid=grid) for galaxy in self.galaxies])
+        return sum(
+            [galaxy.potential_2d_from(grid=grid, xp=xp) for galaxy in self.galaxies]
+        )
 
     @aa.grid_dec.to_array
-    def time_delays_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
+    def time_delays_from(self, grid: aa.type.Grid2DLike, xp=np) -> aa.Array2D:
         """
         Returns the gravitational lensing time delay in days, for a grid of 2D (y, x) coordinates.
 
@@ -773,6 +787,7 @@ def time_delays_from(self, grid: aa.type.Grid2DLike) -> aa.Array2D:
         return tracer_util.time_delays_from(
             galaxies=ag.Galaxies(self.galaxies_ascending_redshift),
             grid=grid,
+            xp=np,
             cosmology=self.cosmology,
         )
 
diff --git a/autolens/lens/tracer_util.py b/autolens/lens/tracer_util.py
index 6cbdbdac9..9fcca5332 100644
--- a/autolens/lens/tracer_util.py
+++ b/autolens/lens/tracer_util.py
@@ -89,6 +89,7 @@ def traced_grid_2d_list_from(
     grid: aa.type.Grid2DLike,
     cosmology: ag.cosmo.LensingCosmology = None,
     plane_index_limit: int = Optional[None],
+    xp=np,
 ):
     """
     Returns a ray-traced grid of 2D Cartesian (y,x) coordinates which accounts for multi-plane ray-tracing.
@@ -165,7 +166,7 @@ def traced_grid_2d_list_from(
                 return traced_grid_list
 
         deflections_yx_2d = sum(
-            map(lambda g: g.deflections_yx_2d_from(grid=scaled_grid), galaxies)
+            map(lambda g: g.deflections_yx_2d_from(grid=scaled_grid, xp=xp), galaxies)
         )
 
         traced_deflection_list.append(deflections_yx_2d)
@@ -257,6 +258,7 @@ def grid_2d_at_redshift_from(
 def time_delays_from(
     galaxies: List[ag.Galaxy],
     grid: aa.type.Grid2DLike,
+    xp=np,
     cosmology: ag.cosmo.LensingCosmology = None,
 ) -> aa.type.Grid2DLike:
     """
diff --git a/autolens/mock.py b/autolens/mock.py
index 6d8361697..bcf8c1499 100644
--- a/autolens/mock.py
+++ b/autolens/mock.py
@@ -1,6 +1,5 @@
-import jax.numpy as jnp
+import numpy as np
 
-from autoconf.jax_wrapper import register_pytree_node_class
 from autofit.mock import *  # noqa
 from autoarray.mock import *  # noqa
 from autogalaxy.mock import *  # noqa
@@ -17,13 +16,13 @@ class NullTracer(Tracer):
     def __init__(self):
         super().__init__([])
 
-    def deflections_yx_2d_from(self, grid):
-        return jnp.zeros_like(grid.array)
+    def deflections_yx_2d_from(self, grid, xp=np):
+        return xp.zeros_like(grid.array)
 
-    def deflections_between_planes_from(self, grid, plane_i=0, plane_j=-1):
-        return jnp.zeros_like(grid.array)
+    def deflections_between_planes_from(self, grid, xp=np, plane_i=0, plane_j=-1):
+        return xp.zeros_like(grid.array)
 
     def magnification_2d_via_hessian_from(
-        self, grid, buffer: float = 0.01, deflections_func=None
+        self, grid, xp=np, buffer: float = 0.01, deflections_func=None
     ) -> aa.ArrayIrregular:
-        return aa.ArrayIrregular(values=jnp.ones(grid.shape[0]))
+        return aa.ArrayIrregular(values=xp.ones(grid.shape[0]))
diff --git a/autolens/plot/__init__.py b/autolens/plot/__init__.py
index 00ccff679..9e66da96d 100644
--- a/autolens/plot/__init__.py
+++ b/autolens/plot/__init__.py
@@ -72,11 +72,8 @@
 
 from autogalaxy.profiles.plot.basis_plotters import BasisPlotter
 from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePlotter
-from autogalaxy.profiles.plot.light_profile_plotters import LightProfilePDFPlotter
 from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePlotter
-from autogalaxy.profiles.plot.mass_profile_plotters import MassProfilePDFPlotter
 from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPlotter
-from autogalaxy.galaxy.plot.galaxy_plotters import GalaxyPDFPlotter
 from autogalaxy.quantity.plot.fit_quantity_plotters import FitQuantityPlotter
 
 from autogalaxy.imaging.plot.fit_imaging_plotters import FitImagingPlotter
diff --git a/autolens/point/fit/abstract.py b/autolens/point/fit/abstract.py
index 2b5ee9a4d..d987b4e23 100644
--- a/autolens/point/fit/abstract.py
+++ b/autolens/point/fit/abstract.py
@@ -1,5 +1,6 @@
 from abc import ABC
 from functools import partial
+import numpy as np
 from typing import Optional, Tuple
 
 import autoarray as aa
@@ -20,6 +21,7 @@ def __init__(
         tracer: Tracer,
         solver: PointSolver,
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Abstract class to fit a point source dataset using a `Tracer` object, including different components
@@ -70,6 +72,8 @@ def __init__(
                 f"in the tracer (make sure your tracer's point source name is the same the dataset name."
             )
 
+        self._xp = xp
+
     @property
     def data(self):
         return self._data
@@ -105,6 +109,7 @@ def deflections_func(self):
 
             return partial(
                 self.tracer.deflections_between_planes_from,
+                xp=self._xp,
                 plane_i=0,
                 plane_j=upper_plane_index,
             )
diff --git a/autolens/point/fit/dataset.py b/autolens/point/fit/dataset.py
index 59ff87902..2b276cbc8 100644
--- a/autolens/point/fit/dataset.py
+++ b/autolens/point/fit/dataset.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from autolens.point.dataset import PointDataset
 from autolens.point.solver import PointSolver
 from autolens.point.fit.fluxes import FitFluxes
@@ -15,6 +17,7 @@ def __init__(
         tracer: Tracer,
         solver: PointSolver,
         fit_positions_cls=FitPositionsImagePair,
+        xp=np,
     ):
         """
         Fits a point source dataset using a `Tracer` object, where the following components of the point source data
@@ -86,6 +89,7 @@ def __init__(
                 tracer=tracer,
                 solver=solver,
                 profile=profile,
+                xp=xp,
             )
         except exc.PointExtractionException:
             self.positions = None
@@ -98,6 +102,7 @@ def __init__(
                     noise_map=dataset.fluxes_noise_map,
                     positions=dataset.positions,
                     tracer=tracer,
+                    xp=xp,
                 )
             else:
                 self.flux = None
@@ -113,12 +118,15 @@ def __init__(
                     noise_map=dataset.time_delays_noise_map,
                     positions=dataset.positions,
                     tracer=tracer,
+                    xp=xp,
                 )
             else:
                 self.time_delays = None
         except exc.PointExtractionException:
             self.time_delays = None
 
+        self._xp = xp
+
     @property
     def model_obj(self):
         return self.tracer
diff --git a/autolens/point/fit/fluxes.py b/autolens/point/fit/fluxes.py
index 898fe6f87..9be71dc07 100644
--- a/autolens/point/fit/fluxes.py
+++ b/autolens/point/fit/fluxes.py
@@ -1,4 +1,4 @@
-import jax.numpy as jnp
+import numpy as np
 from typing import Optional
 
 import autoarray as aa
@@ -19,6 +19,7 @@ def __init__(
         positions: aa.Grid2DIrregular,
         tracer: Tracer,
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Fits the fluxes of a a point source dataset using a `Tracer` object, where every model flux of the point-source
@@ -84,6 +85,7 @@ def __init__(
             tracer=tracer,
             solver=None,
             profile=profile,
+            xp=xp,
         )
 
         if not hasattr(self.profile, "flux"):
@@ -102,7 +104,7 @@ def model_data(self):
         are used.
         """
         return aa.ArrayIrregular(
-            values=jnp.array(
+            values=self._xp.array(
                 [
                     magnification * self.profile.flux
                     for magnification in self.magnifications_at_positions
diff --git a/autolens/point/fit/positions/abstract.py b/autolens/point/fit/positions/abstract.py
index 5b3d176a3..66c3264dd 100644
--- a/autolens/point/fit/positions/abstract.py
+++ b/autolens/point/fit/positions/abstract.py
@@ -1,4 +1,5 @@
 from abc import ABC
+import numpy as np
 from typing import Optional
 
 import autoarray as aa
@@ -18,6 +19,7 @@ def __init__(
         tracer: Tracer,
         solver: PointSolver,
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Abstract class to fit the positions of a a point source dataset using a `Tracer` object, where the specific
@@ -65,6 +67,7 @@ def __init__(
             tracer=tracer,
             solver=solver,
             profile=profile,
+            xp=xp,
         )
 
     @property
diff --git a/autolens/point/fit/positions/image/abstract.py b/autolens/point/fit/positions/image/abstract.py
index 9a0597d62..8b9abb161 100644
--- a/autolens/point/fit/positions/image/abstract.py
+++ b/autolens/point/fit/positions/image/abstract.py
@@ -19,6 +19,7 @@ def __init__(
         tracer: Tracer,
         solver: PointSolver,
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Abstract class to fit the positions of a point source dataset using a `Tracer` object with an image-plane
@@ -70,6 +71,7 @@ def __init__(
             tracer=tracer,
             solver=solver,
             profile=profile,
+            xp=xp,
         )
 
     @staticmethod
diff --git a/autolens/point/fit/positions/image/pair_all.py b/autolens/point/fit/positions/image/pair_all.py
index d93e6bbb4..33e2c3858 100644
--- a/autolens/point/fit/positions/image/pair_all.py
+++ b/autolens/point/fit/positions/image/pair_all.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 
 from autolens.point.fit.positions.image.abstract import AbstractFitPositionsImagePair
@@ -84,7 +83,7 @@ def log_p(
             The log probability of the model coordinate explaining the observed coordinate.
         """
         chi2 = self.square_distance(data_position, model_position) / sigma**2
-        return -jnp.log(jnp.sqrt(2 * jnp.pi * sigma**2)) - 0.5 * chi2
+        return -self._xp.log(self._xp.sqrt(2 * self._xp.pi * sigma**2)) - 0.5 * chi2
 
     def all_permutations_log_likelihoods(self) -> np.ndarray:
         """
@@ -103,13 +102,13 @@ def all_permutations_log_likelihoods(self) -> np.ndarray:
 
         model_data = self.model_data.array
 
-        return jnp.array(
+        return self._xp.array(
             [
-                jnp.log(
-                    jnp.sum(
-                        jnp.array(
+                self._xp.log(
+                    self._xp.sum(
+                        self._xp.array(
                             [
-                                jnp.exp(
+                                self._xp.exp(
                                     self.log_p(
                                         data_position,
                                         model_position,
@@ -143,12 +142,13 @@ def chi_squared(self) -> float:
 
         This is every way in which the coordinates generated by the model can explain the observed coordinates.
         """
-        n_non_nan_model_positions = jnp.count_nonzero(
-            jnp.isfinite(
+        n_non_nan_model_positions = self._xp.count_nonzero(
+            self._xp.isfinite(
                 self.model_data.array,
             ).any(axis=1)
         )
         n_permutations = n_non_nan_model_positions ** len(self.data)
         return -2.0 * (
-            -jnp.log(n_permutations) + jnp.sum(self.all_permutations_log_likelihoods())
+            -self._xp.log(n_permutations)
+            + self._xp.sum(self.all_permutations_log_likelihoods())
         )
diff --git a/autolens/point/fit/positions/image/pair_repeat.py b/autolens/point/fit/positions/image/pair_repeat.py
index 7d091f647..e8d247531 100644
--- a/autolens/point/fit/positions/image/pair_repeat.py
+++ b/autolens/point/fit/positions/image/pair_repeat.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import autoarray as aa
 
 from autolens.point.fit.positions.image.abstract import AbstractFitPositionsImagePair
@@ -62,6 +61,6 @@ def residual_map(self) -> aa.ArrayIrregular:
                 self.square_distance(model_position, position)
                 for model_position in self.model_data.array
             ]
-            residual_map.append(jnp.sqrt(jnp.min(jnp.array(distances))))
+            residual_map.append(self._xp.sqrt(self._xp.min(self._xp.array(distances))))
 
-        return aa.ArrayIrregular(values=jnp.array(residual_map))
+        return aa.ArrayIrregular(values=self._xp.array(residual_map))
diff --git a/autolens/point/fit/positions/source/separations.py b/autolens/point/fit/positions/source/separations.py
index 3a4529d7a..a06280069 100644
--- a/autolens/point/fit/positions/source/separations.py
+++ b/autolens/point/fit/positions/source/separations.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional
 
@@ -19,6 +18,7 @@ def __init__(
         tracer: Tracer,
         solver: Optional[PointSolver],
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Fits the positions of a a point source dataset using a `Tracer` object with a source-plane chi-squared based on
@@ -79,6 +79,7 @@ def __init__(
             tracer=tracer,
             solver=solver,
             profile=profile,
+            xp=xp,
         )
 
     @property
@@ -95,7 +96,7 @@ def model_data(self) -> aa.Grid2DIrregular:
             deflections = self.tracer.deflections_yx_2d_from(grid=self.data)
         else:
             deflections = self.tracer.deflections_between_planes_from(
-                grid=self.data, plane_i=0, plane_j=self.plane_index
+                grid=self.data, xp=self._xp, plane_i=0, plane_j=self.plane_index
             )
 
         return self.data.grid_2d_via_deflection_grid_from(deflection_grid=deflections)
@@ -126,8 +127,8 @@ def noise_normalization(self) -> float:
         """
         Returns the normalization of the noise-map, which is the sum of the noise-map values squared.
         """
-        return jnp.sum(
-            jnp.log(
+        return self._xp.sum(
+            self._xp.log(
                 2
                 * np.pi
                 * (
diff --git a/autolens/point/fit/times_delays.py b/autolens/point/fit/times_delays.py
index 23a2cae55..598abf01f 100644
--- a/autolens/point/fit/times_delays.py
+++ b/autolens/point/fit/times_delays.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import numpy as np
 from typing import Optional
 
@@ -18,6 +17,7 @@ def __init__(
         positions: aa.Grid2DIrregular,
         tracer: Tracer,
         profile: Optional[ag.ps.Point] = None,
+        xp=np,
     ):
         """
         Fits the time delays of a point source dataset using a `Tracer` object,
@@ -63,6 +63,7 @@ def __init__(
             tracer=tracer,
             solver=None,
             profile=profile,
+            xp=xp,
         )
 
     @property
@@ -74,7 +75,7 @@ def model_data(self) -> aa.ArrayIrregular:
         delay have a value of zero. However, this subtraction is performed in the `residual_map` property, in order
         to ensure the residuals are computed relative to the shorter time delay.
         """
-        return self.tracer.time_delays_from(grid=self.positions)
+        return self.tracer.time_delays_from(grid=self.positions, xp=self._xp)
 
     @property
     def model_time_delays(self) -> aa.ArrayIrregular:
@@ -90,8 +91,8 @@ def residual_map(self) -> aa.ArrayIrregular:
         from the dataset time delays and model time delays before the subtraction.
         """
 
-        data = self.data - jnp.min(self.data.array)
-        model_data = self.model_data - jnp.min(self.model_data.array)
+        data = self.data - self._xp.min(self.data.array)
+        model_data = self.model_data - self._xp.min(self.model_data.array)
 
         residual_map = aa.util.fit.residual_map_from(data=data, model_data=model_data)
         return aa.ArrayIrregular(values=residual_map)
diff --git a/autolens/point/max_separation.py b/autolens/point/max_separation.py
index a83af2e1e..6b0fbb55c 100644
--- a/autolens/point/max_separation.py
+++ b/autolens/point/max_separation.py
@@ -1,3 +1,4 @@
+import numpy as np
 from typing import Optional
 
 import autoarray as aa
@@ -12,6 +13,7 @@ def __init__(
         noise_map: Optional[aa.ArrayIrregular],
         tracer: Tracer,
         plane_redshift: float = Optional[None],
+        xp=np,
     ):
         """
         Given a positions dataset, which is a list of positions with names that associated them to model source
@@ -43,7 +45,7 @@ def __init__(
             plane_index = -1
 
         self.plane_positions = aa.Grid2DIrregular(
-            values=tracer.traced_grid_2d_list_from(grid=data)[plane_index]
+            values=tracer.traced_grid_2d_list_from(grid=data, xp=xp)[plane_index], xp=xp
         )
 
     @property
diff --git a/autolens/point/solver/point_solver.py b/autolens/point/solver/point_solver.py
index 82e7067f5..9cdcbbd1e 100644
--- a/autolens/point/solver/point_solver.py
+++ b/autolens/point/solver/point_solver.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import logging
 from typing import Tuple, Optional
 
@@ -46,6 +45,8 @@ def solve(
         -------
         A list of image plane coordinates that are traced to the source plane coordinate.
         """
+        import jax.numpy as jnp
+
         kept_triangles = super().solve_triangles(
             tracer=tracer,
             shape=Point(*source_plane_coordinate),
diff --git a/autolens/point/solver/shape_solver.py b/autolens/point/solver/shape_solver.py
index 2b0980f60..506d2beaa 100644
--- a/autolens/point/solver/shape_solver.py
+++ b/autolens/point/solver/shape_solver.py
@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 import logging
 import math
 
@@ -254,6 +253,8 @@ def _filter_low_magnification(
         -------
         The points with an absolute magnification above the threshold.
         """
+        import jax.numpy as jnp
+
         points = jnp.array(points)
         magnifications = tracer.magnification_2d_via_hessian_from(
             grid=aa.Grid2DIrregular(points).array,
diff --git a/docs/api/pixelization.rst b/docs/api/pixelization.rst
index 051441174..7a8485cab 100644
--- a/docs/api/pixelization.rst
+++ b/docs/api/pixelization.rst
@@ -46,7 +46,7 @@ Mesh [ag.mesh]
    :template: custom-class-template.rst
    :recursive:
 
-   Rectangular
+   RectangularMagnification
    Delaunay
    Voronoi
    Voronoi
diff --git a/docs/api/plot.rst b/docs/api/plot.rst
index bef26ff77..acac39c44 100644
--- a/docs/api/plot.rst
+++ b/docs/api/plot.rst
@@ -32,7 +32,6 @@ Create figures and subplots showing quantities of standard **PyAutoLens** object
     ImagingPlotter
     InterferometerPlotter
     LightProfilePlotter
-    LightProfilePDFPlotter
     GalaxyPlotter
     FitImagingPlotter
     FitInterferometerPlotter
diff --git a/test_autolens/config/general.yaml b/test_autolens/config/general.yaml
index f13037268..9bbac4100 100644
--- a/test_autolens/config/general.yaml
+++ b/test_autolens/config/general.yaml
@@ -17,8 +17,6 @@ inversion:
   use_positive_only_solver: false    # If True, inversion's use a positive-only linear algebra solver by default, which is slower but prevents unphysical negative values in the reconstructed solutuion.
   no_regularization_add_to_curvature_diag_value: 1.0e-8 # The default value added to the curvature matrix's diagonal when regularization is not applied to a linear object, which prevents inversion's failing due to the matrix being singular.
   positive_only_uses_p_initial: false  # If True, the positive-only solver of an inversion's uses an initial guess of the reconstructed data's values as which values should be positive, speeding up the solver.
-model:
-  ignore_prior_limits: true
 numba:
   cache: true
   nopython: true
diff --git a/test_autolens/config/priors/pixelizations.yaml b/test_autolens/config/priors/pixelizations.yaml
index ac760c154..b9eb85521 100644
--- a/test_autolens/config/priors/pixelizations.yaml
+++ b/test_autolens/config/priors/pixelizations.yaml
@@ -50,7 +50,7 @@ delaunay.DelaunayMagnification:
     width_modifier:
       type: Absolute
       value: 8.0
-rectangular.Rectangular:
+rectangular.RectangularMagnification:
   shape_0:
     limits:
       lower: 3.0
diff --git a/test_autolens/lens/test_tracer.py b/test_autolens/lens/test_tracer.py
index 626b0dc19..0c664708a 100644
--- a/test_autolens/lens/test_tracer.py
+++ b/test_autolens/lens/test_tracer.py
@@ -323,7 +323,7 @@ def test__image_2d_via_input_plane_image_from__with_foreground_planes(grid_2d_7x
     )
 
     assert image_via_light_profile[0] == pytest.approx(
-        image_via_input_plane_image[0].array, 1.0e-2
+        image_via_input_plane_image[0], 1.0e-2
     )
 
 
@@ -398,7 +398,7 @@ def test__image_2d_via_input_plane_image_from__with_foreground_planes__multi_pla
     )
 
     assert image_via_light_profile[0] == pytest.approx(
-        image_via_input_plane_image[0].array, 1.0e-2
+        image_via_input_plane_image[0], 1.0e-2
     )
 
     plane_image = g1.image_2d_from(grid=plane_grid)
@@ -411,7 +411,7 @@ def test__image_2d_via_input_plane_image_from__with_foreground_planes__multi_pla
     )
 
     assert image_via_light_profile[0] == pytest.approx(
-        image_via_input_plane_image[0].array, 1.0e-2
+        image_via_input_plane_image[0], 1.0e-2
     )