feat: public register_galaxies_classes, the autolens register_tracer_classes counterpart

[Jammy2211]

Summary

PyAutoLens ships autolens.jax.register_tracer_classes(tracer) — the public,
one-time JAX pytree registration for code paths that do not go through an
Analysis: a user's own @jax.jit that receives the aggregate object as a
traced argument. PyAutoGalaxy had no equivalent public entry point.

Surfaced while writing the __Custom Likelihood Functions__ section of
autogalaxy_workspace/scripts/guides/using_jax.py
(PyAutoLabs/autogalaxy_workspace#181, merged): the guide could document the
Analysis, hand-rolled-inside-jit and Fitness paths, but not passing a
Galaxies as a jit argument, because no public call made it work.

Why the existing helper wasn't enough

autogalaxy/analysis/jax_pytrees.py::register_galaxies_pytree registers only
the Galaxies list subclass, and is private — reached solely from the two
Analysis classes via fit_from. Measured on the installed stack with:

@jax.jit
def log_likelihood(galaxies):
    return ag.FitImaging(dataset=masked_dataset, galaxies=galaxies, xp=jnp).log_likelihood

Registration	Result
none	TypeError: ... problematic value is of type Galaxies ... at path galaxies
register_galaxies_pytree() only	still fails — ... of type Galaxy ... at path galaxies[0]
both steps	works; -270175.0553756637, exactly the eager value

What this adds

• autogalaxy/jax/__init__.py — exports register_galaxies_classes, mirroring autolens/jax/__init__.py.
• autogalaxy/jax/registration.py — register_galaxies_classes(galaxies) -> bool, which calls the existing shared register_galaxies_pytree() and then walks each galaxy registering Galaxy and every profile class. Idempotent; returns False as a silent no-op when JAX is absent.

Nothing else changes: register_galaxies_pytree, both Analysis classes and
all of autolens are untouched. autogalaxy.jax is auto-discovered by
[tool.setuptools.packages.find], so no packaging change is needed.

Two things a reviewer should know

1. Nothing inside autogalaxy calls this. Unlike autolens — where
PointSolver and Simulator invoke register_tracer_classes automatically —
both autogalaxy Analysis classes already register their pytrees inline from
fit_from. This is purely the public entry point for user code. The module
docstring says so explicitly so it isn't later deleted as dead code.

2. The ~60-line recursive walker is duplicated from
autolens/jax/registration.py rather than shared. Deliberate, and the human's
call: it keeps this to one repo and one reversible PR instead of coupling
autolens main to autogalaxy main immediately after the 2026.7.29.2
release. Deduping (autogalaxy would own it; autolens may import autogalaxy) is
recorded as a follow-up, not done here.

Test Plan

• python -m pytest test_autogalaxy/ — 1009 passed, 0 failed
• Three-way probe against this branch: none → fails on Galaxies; register_galaxies_pytree() only → fails on Galaxy; register_galaxies_classes(galaxies) → works and matches the eager log likelihood exactly
• Idempotent — a second call returns True and re-registers nothing
• import autogalaxy does not eagerly import autogalaxy.jax (verified via sys.modules), so the new module adds no import cost
• With import jax forced to raise, register_galaxies_classes returns False and does not raise

Unit tests in this repo are NumPy-only by policy, so the JAX behaviour is
validated by the probe rather than by test_autogalaxy/.

Follow-ups (not in this PR)

• Workspace guide mention — document the jit-argument case and this call in autogalaxy_workspace/scripts/guides/using_jax.py, matching how the autolens guide points at register_tracer_classes. Blocked on this PR merging (library-first).
• Dedupe the walker across autolens/jax/registration.py and autogalaxy/jax/registration.py.
• Verify a neighbouring doc claim: using_jax.py states "The simulator handles pytree registration internally", but autogalaxy/imaging/simulator.py and interferometer/simulator.py contain no pytree registration calls. Either an autoarray base does it, or that sentence is wrong too. Pre-existing text, not introduced here.

Generated by the PyAutoLabs agent workflow.