feat: register FitImaging, DatasetModel, Galaxies as pytrees in AnalysisImaging

Adds `_register_fit_imaging_pytrees` staticmethod to `AnalysisImaging`, called
from `fit_from` under the existing `use_jax` gate. Mirrors autolens's
`AnalysisImaging._register_fit_imaging_pytrees`. Unblocks
`jax.jit(analysis.fit_from)(instance)` returning a `FitImaging` with
`jax.Array` leaves on the autogalaxy imaging path.

Part of PyAutoLabs/autogalaxy_workspace_test#8 (epic #5, task 3/9).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: 2 people

autogalaxy/imaging/model/analysis.py

@@ -141,6 +141,9 @@ def fit_from(self, instance: af.ModelInstance) -> FitImaging:
             The fit of the galaxies to the imaging dataset, which includes the log likelihood.
         """
 
+        if self._use_jax:
+            self._register_fit_imaging_pytrees()
+
         galaxies = self.galaxies_via_instance_from(
             instance=instance,
         )
@@ -158,6 +161,27 @@ def fit_from(self, instance: af.ModelInstance) -> FitImaging:
             xp=self._xp,
         )
 
+    @staticmethod
+    def _register_fit_imaging_pytrees() -> None:
+        """Register every type reachable from a ``FitImaging`` return value
+        so ``jax.jit(fit_from)`` can flatten its output.
+
+        ``dataset``, ``adapt_images`` and ``settings`` are constants per
+        analysis — ride as aux so JAX does not recurse into them. Everything
+        else (``galaxies``, ``dataset_model`` and the autoarray wrappers they
+        carry) is dynamic per fit.
+        """
+        from autoarray.abstract_ndarray import register_instance_pytree
+        from autoarray.dataset.dataset_model import DatasetModel
+        from autogalaxy.galaxy.galaxies import Galaxies
+
+        register_instance_pytree(
+            FitImaging,
+            no_flatten=("dataset", "adapt_images", "settings"),
+        )
+        register_instance_pytree(DatasetModel)
+        register_instance_pytree(Galaxies)
+
     def save_attributes(self, paths: af.DirectoryPaths):
         """
         Before the non-linear search begins, this routine saves attributes of the `Analysis` object to the `files`

test_autogalaxy/imaging/model/test_analysis_imaging.py

@@ -37,3 +37,16 @@ def test__figure_of_merit__matches_correct_fit_given_galaxy_profiles(
     fit = ag.FitImaging(dataset=masked_imaging_7x7, galaxies=galaxies)
 
     assert fit.log_likelihood == fit_figure_of_merit
+
+
+def test__register_fit_imaging_pytrees__registers_fit_galaxies_and_dataset_model():
+    from autoarray.abstract_ndarray import _pytree_registered_classes
+    from autoarray.dataset.dataset_model import DatasetModel
+    from autogalaxy.galaxy.galaxies import Galaxies
+    from autogalaxy.imaging.fit_imaging import FitImaging
+
+    ag.AnalysisImaging._register_fit_imaging_pytrees()
+
+    assert FitImaging in _pytree_registered_classes
+    assert DatasetModel in _pytree_registered_classes
+    assert Galaxies in _pytree_registered_classes