Feature/jax in image dict

[Jammy2211]

This pull request introduces changes to support passing a custom array module (such as numpy or jpy, referenced as xp) throughout the galaxy image and visibilities calculation pipeline in both imaging and interferometer fitting. This makes the code more flexible and better suited for GPU acceleration or alternative array backends. Additionally, the relevant test has been updated for formatting consistency.

Backend flexibility improvements:

• Updated galaxy_image_2d_dict_from in tracer.py to accept an xp parameter (defaulting to np), and ensured all internal calls and loops propagate this parameter, enabling support for different array libraries. [1] [2]
• Modified fit_imaging.py and fit_interferometer.py to pass self._xp as the xp parameter when calling tracer methods for generating galaxy image and model dictionaries, ensuring consistent backend usage throughout the pipeline. [1] [2] [3] [4]

Testing and code style:

• Reformatted assertions in test_fit_interferometer.py for improved readability and consistency.

[Copilot]

Pull request overview

This PR extends the image/visibilities dictionary generation pipeline to consistently propagate a configurable array backend (xp, e.g. NumPy/JAX) through Tracer and the FitImaging / FitInterferometer model-generation paths, enabling alternative array backends (including GPU-friendly ones).

Changes:

• Added an xp parameter to Tracer.galaxy_image_2d_dict_from and propagated it through traced-grid and per-galaxy image evaluation.
• Updated imaging/interferometer fit classes to pass self._xp into tracer dictionary generation.
• Reformatted a couple of assertions in the interferometer fit test.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
autolens/lens/tracer.py	Adds xp support to galaxy_image_2d_dict_from and propagates it internally.
autolens/imaging/fit_imaging.py	Passes self._xp into tracer galaxy image/model dictionary generation.
autolens/interferometer/fit_interferometer.py	Passes self._xp into tracer galaxy image and visibilities dictionary generation.
test_autolens/interferometer/test_fit_interferometer.py	Assertion formatting adjustments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Adding xp as the second positional parameter changes the positional call signature (e.g. existing user code like galaxy_image_2d_dict_from(grid, True) would now pass True as xp). To preserve backward compatibility, keep operated_only as the second parameter and make xp keyword-only (e.g. by moving xp after operated_only or introducing a * before xp).

Suggested change

	self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
	self,
	grid: aa.type.Grid2DLike,
	operated_only: Optional[bool] = None,
	*,
	xp=np,