docs: registration is the caller's job — correct stale auto-call claims in autolens/jax - #665
Merged
Merged
Conversation
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 <[email protected]> Claude-Session: https://claude.ai/code/session_01CWqjHGXUut25TEB8octU8H
Collaborator
Author
|
Workspace PRs: PyAutoLabs/autolens_workspace#380, PyAutoLabs/autogalaxy_workspace#183 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The
autolens/jax/registration.pymodule docstring claimsregister_tracer_classes"is called automatically byPointSolver(use_jax=True).solve(tracer, ...)on the first invocation and bySimulator(use_jax=True).via_tracer_from(tracer, ...)in PyAutoLens once Phase 2ships the Simulator changes."
Neither is true.
grep -rn register_tracer_classes autolens/finds only thedefinition and the re-export — nothing in the library calls it. And it cannot
be called automatically:
PointSolver.solve_trianglesalready carries the reasonat its own call site (
autolens/point/solver/point_solver.py:102-109):So the docstring promises behaviour that is not merely unshipped but
unimplementable as worded, and it sends readers looking for automatic
registration that will never arrive.
This corrects the docstring to state that calling it is the user's
responsibility, gives the trace-time reason, and cross-references the new
autogalaxy counterpart
autogalaxy.jax.register_galaxies_classes(added inPyAutoGalaxy#537).
Docstring only — zero behaviour change, zero API surface.
Context
Found while correcting the same false claim where it had propagated into the
workspaces: both
scripts/guides/using_jax.pyfiles and the__JAX Variant__sections of several
simulator.pyscripts assert "the simulator handles pytreeregistration internally". Those are fixed in the companion workspace PRs;
tracking issue PyAutoLabs/autolens_workspace#379.
Separately, the jitted simulator path fails even with correct registration, on
un-threaded
xpsites in autoarray — that is a real library bug, filedseparately, and not addressed here.
Test Plan
python -m pytest test_autolens/— 488 passed, 0 failedGenerated by the PyAutoLabs agent workflow.