From 4b2decdc7ccc186076e16fe46d82a1236900ca8e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 30 Jul 2026 10:06:32 +0100 Subject: [PATCH] =?UTF-8?q?docs:=20registration=20is=20the=20caller's=20jo?= =?UTF-8?q?b=20=E2=80=94=20correct=20stale=20auto-call=20claims?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The autolens/jax/registration.py module docstring claimed register_tracer_classes "is called automatically by PointSolver(use_jax=True).solve(tracer, ...) on the first invocation and by Simulator(use_jax=True).via_tracer_from(tracer, ...) ... once Phase 2 ships the Simulator changes". Neither is true. grep finds no caller anywhere in autolens. And it cannot be done automatically: PointSolver.solve_triangles already carries the reason at its own call site (point_solver.py:102-109) — JAX flattens a jitted function's arguments at trace time, before entering the callee, so registration must run before the first jitted call. Corrected to state that calling it is the user's responsibility, with the trace-time reason, and to cross-reference the new autogalaxy counterpart autogalaxy.jax.register_galaxies_classes (PyAutoGalaxy#537). Docstring only — no behaviour change, no API surface. Refs PyAutoLabs/autolens_workspace#379 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CWqjHGXUut25TEB8octU8H --- autolens/jax/registration.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/autolens/jax/registration.py b/autolens/jax/registration.py index 78573fa01..a71a4ab33 100644 --- a/autolens/jax/registration.py +++ b/autolens/jax/registration.py @@ -7,10 +7,16 @@ This module is the counterpart of ``AnalysisImaging._register_fit_imaging_pytrees`` for code paths that do not go through ``Analysis`` (point-source solving, -custom forward models, hand-built simulators). It is called automatically -by ``PointSolver(use_jax=True).solve(tracer, ...)`` on the first invocation -and by ``Simulator(use_jax=True).via_tracer_from(tracer, ...)`` in PyAutoLens -once Phase 2 ships the Simulator changes. +custom forward models, hand-built simulators). + +**Calling it is the user's responsibility.** Nothing in the library calls it +for you, and nothing can: JAX flattens a jitted function's arguments at trace +time, i.e. *before* entering the callee, so a ``solve()`` or +``via_tracer_from()`` that registered internally would already be too late. +``PointSolver.solve_triangles`` carries the same note at its own call site. +Register once, before the first ``@jax.jit`` invocation. + +The autogalaxy counterpart is ``autogalaxy.jax.register_galaxies_classes``. Mirrors PyAutoFit's ``autofit/jax/pytrees.py`` layout. Idempotent: re-registration of a class is a silent no-op.