Skip to content

docs: registration is the caller's job — correct stale auto-call claims in autolens/jax - #665

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/correct-simulator-jax-claims
Jul 30, 2026
Merged

docs: registration is the caller's job — correct stale auto-call claims in autolens/jax#665
Jammy2211 merged 1 commit into
mainfrom
feature/correct-simulator-jax-claims

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

The autolens/jax/registration.py module docstring claims
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, ...) in PyAutoLens once Phase 2
ships the Simulator changes."

Neither is true. grep -rn register_tracer_classes autolens/ finds only the
definition and the re-export — nothing in the library calls it. And it cannot
be called automatically: PointSolver.solve_triangles already carries the reason
at its own call site (autolens/point/solver/point_solver.py:102-109):

Auto-registering inside solve() doesn't help because JAX flattens function
arguments at trace time — before entering this method — so registration must
run before the first jitted call.

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 in
PyAutoGalaxy#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.py files and the __JAX Variant__
sections of several simulator.py scripts assert "the simulator handles pytree
registration 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 xp sites in autoarray — that is a real library bug, filed
separately, and not addressed here.

Test Plan

  • python -m pytest test_autolens/488 passed, 0 failed
  • Verified by grep that no caller exists, so the removed claim describes nothing real

Generated by the PyAutoLabs agent workflow.

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
@Jammy2211

Copy link
Copy Markdown
Collaborator Author

@Jammy2211
Jammy2211 merged commit d53bfcc into main Jul 30, 2026
5 checks passed
@Jammy2211
Jammy2211 deleted the feature/correct-simulator-jax-claims branch July 30, 2026 09:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant