Register FitInterferometer return value as a JAX pytree

Adds ``_register_fit_interferometer_pytrees`` to ``AnalysisInterferometer``,
mirroring the imaging-side ``_register_fit_imaging_pytrees``. Called at the
top of ``fit_from`` when ``use_jax=True`` so ``jax.jit(analysis.fit_from)``
can flatten the returned ``FitInterferometer``.

Registers:
- ``FitInterferometer`` with dataset / adapt_images / settings as aux
- ``Tracer`` with cosmology as aux (already covered by the imaging path
  but harmless to re-register under the interferometer site)
- ``DatasetModel`` — always present on ``FitDataset`` subclasses even
  when ``fit_from`` doesn't pass one explicitly (``FitDataset.__init__``
  falls back to ``DatasetModel()``)

Unlocks the Path A PoC in
``autolens_workspace_test/scripts/jax_likelihood_functions/interferometer/mge_pytree.py``
for the MGE-source interferometer variant.

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

Author: 2 people

autolens/interferometer/model/analysis.py

@@ -173,6 +173,9 @@ def fit_from(self, instance: af.ModelInstance) -> FitInterferometer:
             The fit of the plane to the interferometer dataset, which includes the log likelihood.
         """
 
+        if self._use_jax:
+            self._register_fit_interferometer_pytrees()
+
         tracer = self.tracer_via_instance_from(
             instance=instance,
         )
@@ -187,6 +190,27 @@ def fit_from(self, instance: af.ModelInstance) -> FitInterferometer:
             xp=self._xp,
         )
 
+    @staticmethod
+    def _register_fit_interferometer_pytrees() -> None:
+        """Register every type reachable from a ``FitInterferometer`` 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 (``tracer`` and the autoarray wrappers it carries) is dynamic
+        per fit.
+        """
+        from autoarray.abstract_ndarray import register_instance_pytree
+        from autoarray.dataset.dataset_model import DatasetModel  # fit-interferometer-pytree-mge
+        from autolens.lens.tracer import Tracer
+
+        register_instance_pytree(
+            FitInterferometer,
+            no_flatten=("dataset", "adapt_images", "settings"),
+        )
+        register_instance_pytree(Tracer, no_flatten=("cosmology",))
+        register_instance_pytree(DatasetModel)  # fit-interferometer-pytree-mge
+
     def save_attributes(self, paths: af.DirectoryPaths):
         """
          Before the model-fit begins, this routine saves attributes of the `Analysis` object to the `files` folder