docs(agents): pytree registration classification heuristic - #491
Merged
Merged
Conversation
Adds §6 to docs/agents/jax_and_decorators.md, beside the §4 Pattern-2 material that already documents register_instance_pytree: how to classify an offending type's attributes when a jit trace fails on an unregistered class (all-array -> no_flatten=(); the known-aux name set — cosmology, settings, dataset, psf, mask, caches, scipy.spatial.*, Transformer*, PointSolver* —> aux; a callable attribute decided deliberately rather than guessed), where to register it, and the iterative re-trace habit with its stop condition. Carries the <variant>_pytree.py round-trip assertion used by the parity scripts in autolens_workspace_test/scripts/jax_likelihood_functions/, and cross-references §5: the round trip proves the types flatten, it does not prove xp is threaded — that still needs fitness._vmap(parameters). Rescued from PyAutoBrain's register_and_iterate skill, retired in that repo this session; the loop mechanics around it were already /run_queue's. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01G1SYyCvBWevmYu3x74N8N2
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
1 participant
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.
Adds §6 — Registering a new type: classification heuristic to
docs/agents/jax_and_decorators.md, directly after the §4 Pattern-2 material that already documentsregister_instance_pytree.Why
The heuristic lived in
PyAutoBrain/skills/register_and_iterate/reference.md. That skill was superseded by/run_queueon 2026-07-08 and has not been invoked since, so it was retired this session — but the pytree knowledge inside it is still current and belongs in the repo that owns the API.What it documents
no_flatten=(); the known-aux name set (cosmology,settings,dataset,psf,mask, caches,scipy.spatial.*,Transformer*,PointSolver*) → aux; a callable attribute decided deliberately rather than guessed._register_fit_imaging_pytrees) and the iterative re-trace habit, with the stop condition — no progress across a few passes means the type holds state that cannot be flattened, which is a design question, not more registrations.<variant>_pytree.pyround-trip assertion used by the parity scripts underautolens_workspace_test/scripts/jax_likelihood_functions/.Cross-referenced to §5 so it does not contradict the doc it joins: the round trip proves the types flatten, it does not prove
xpis threaded — that still needsfitness._vmap(parameters).Scope
Documentation only — no source, no tests, no API change.
🤖 Generated with Claude Code
https://claude.ai/code/session_01G1SYyCvBWevmYu3x74N8N2