fix: JAX pytree registers leaf types (int) → _partition vars() crash

[Jammy2211]

Overview

JAX-fitting a shapelet model crashes when autofit flattens the model instance into a pytree. Surfaced by the 2026-07-13 release-validation Stage-3 run (PyAutoHeart workspace-validation run_id=29266305445, TestPyPI 2026.7.13.1.dev65501) — autogalaxy imaging scripts/imaging/features/shapelets/modeling.py FAILs under the JAX release profile in jax.vmap(jax.jit(self.call)) (Fitness._vmap).

Root cause (diagnosed — differs from the prompt's initial hypothesis)

The prompt guessed a __slots__/array-like component. The real cause is erroneous registration of a leaf type:

• af.Model(ShapeletPolar) auto-wraps its int-typed constructor args n/m as af.Model(int) sub-models.
• register_model._walk (autofit/jax/pytrees.py) registers every Model.cls it walks as a JAX pytree node — including builtins.int.
• Once int is a registered pytree node, JAX calls the instance-flatten _partition(5) on every int leaf (n, m, id, pytree_token), which does vars(5) → TypeError: vars() argument must have __dict__ attribute.

int (like float/str/bool) is a native JAX leaf — it must never be registered as a pytree container. Reproduced in pure autofit (no autogalaxy) via af.Model(int) as a sub-model → int in _REGISTERED_INSTANCE_CLASSES → tree_flatten crashes. The Basis/ShapeletPolar/Galaxy instances themselves all have __dict__ and flatten fine.

Fix (at the producer, per no-silent-guards)

Gate registration in _walk: only register cls whose instances carry a __dict__ — any("__dict__" in vars(k) for k in cls.__mro__). This cleanly excludes int/float/str/bool/NoneType/tuple (native leaves, left opaque) and any purely __slots__ class (correctly treated as an opaque leaf, matching the prompt's "treat as opaque leaf as appropriate"), while still registering real model-component classes. No silent fallback inside _partition — the erroneous registration simply never happens.

Plan

• autofit/jax/pytrees.py: add _instances_carry_dict(cls) helper; in register_model._walk, skip the register_pytree_node block for classes that fail it (leave classifier/ctor-arg bookkeeping intact). Comment the reasoning.
• Regression assertion in autofit_workspace_test/scripts/jax_assertions/ (JAX unit tests were moved out of test_autofit/ in test: delete jax-using unit tests (moved to autofit_workspace_test) #1247): build a model with an af.Model(int) sub-model, register_model, assert int/float are NOT registered and jax.tree_util.tree_flatten(instance) round-trips.
• Verify the crash is gone and the shapelet-shaped tree flattens.

Detailed implementation plan

Work Classification

Library + Workspace (library-first).

Affected Repositories

• PyAutoLabs/PyAutoFit (primary — the _walk registration guard)
• PyAutoLabs/autofit_workspace_test (JAX regression assertion)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./autofit_workspace_test	main	clean

Suggested branch: feature/jax-pytree-leaf-registration
Worktree root: ~/Code/PyAutoLabs-wt/jax-pytree-leaf-registration/

Implementation Steps

1. PyAutoFit/autofit/jax/pytrees.py — helper _instances_carry_dict(cls) and gate the register_pytree_node block in _walk on it.
2. autofit_workspace_test/scripts/jax_assertions/pytree_leaf_registration.py (new) — reproduce af.Model(int) sub-model; assert no leaf-type registration + tree_flatten round-trip.
3. Verify against the reproduction (crash gone; int not registered; flatten OK).

Traceback (as reported)

autofit/non_linear/fitness.py:486   func = jax.vmap(jax.jit(self.call))
autofit/jax/pytrees.py:168 (flatten) -> :152 (_partition)
    for name, value in vars(instance).items():
TypeError: vars() argument must have __dict__ attribute

Autonomy

--auto, effective level supervised (bug work-type cap). Plan written here; run parks at ship sign-off with a question rather than opening the PR. Merge/close remain human.

Original Prompt

Click to expand starting prompt

JAX pytree _partition crashes on instances without dict (vars() TypeError) during vmap/jit

Type: bug
Target: autofit
Repos:

• PyAutoFit
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: formalised

JAX-fitting a shapelet model crashes when autofit flattens the model instance into a
pytree. Surfaced by the 2026-07-13 release-validation Stage-3 run (PyAutoHeart
workspace-validation run_id=29266305445, TestPyPI 2026.7.13.1.dev65501) — autogalaxy
imaging scripts/imaging/features/shapelets/modeling.py FAIL under the release env profile
(JAX enabled).

Traceback:

scripts/imaging/features/shapelets/modeling.py:424   result = search.fit(model=model, analysis=analysis)
autofit/non_linear/search/nest/nautilus/search.py:189   fitness = Fitness(...)
autofit/non_linear/fitness.py:134   self._call = self._vmap
autofit/non_linear/fitness.py:486   func = jax.vmap(jax.jit(self.call))
autofit/jax/pytrees.py:168 (flatten) -> :152 (_partition)
    for name, value in vars(instance).items():
TypeError: vars() argument must have __dict__ attribute

Root cause — autofit/jax/pytrees.py:152 (_partition) assumes every model instance /
component exposes __dict__; it calls vars(instance) unconditionally. A shapelet model
carries a component whose instance has no __dict__ (e.g. __slots__-based or an
array-like leaf), so vars() raises. This blocks jax.vmap(jax.jit(self.call)) in
Fitness._vmap, i.e. any JAX-accelerated fit whose model contains such a component.

Only exercised under the JAX/release profile (the curated smoke subset does not JAX-fit
shapelets), so the nightly release would not catch it — but it breaks JAX shapelet fitting
for users. Fix _partition to handle no-__dict__ instances (fall back to
__slots__ traversal / treat as an opaque leaf as appropriate) rather than assuming
vars(); add a JAX pytree-flatten unit test covering a slotted/array-like component.
Do NOT paper over it with a silent guard — see [[feedback_no_silent_guards]]. Cross-check
against the existing pytree PoC work ([[register_and_iterate]] history) for the intended
partition contract.

[Jammy2211]

⏸️ Parked at ship sign-off — supervised --auto (needs your go-ahead to open the PR)

Implementation is done and verified; under the supervised bug cap the run stops here for your review before commit/push/PR. Nothing committed or pushed. Merge/close stay human.

Diagnosis note (differs from the prompt's hypothesis)

The prompt guessed a __slots__/array-like component. The true root cause is erroneous registration of the leaf type int: af.Model(ShapeletPolar) wraps its int args n/m as af.Model(int) sub-models, and register_model._walk registered builtins.int as a pytree node → JAX then ran the instance-flatten _partition(5) on every int leaf → vars(5) crash. Confirmed by walking the tree (int ∈ _REGISTERED_INSTANCE_CLASSES) and reproduced in pure autofit via af.Model(int).

What changed (worktree ~/Code/PyAutoLabs-wt/jax-pytree-leaf-registration, branch feature/jax-pytree-leaf-registration)

PyAutoFit (autofit/jax/pytrees.py, +29/-1)

• New helper _instances_carry_dict(cls) = any("__dict__" in vars(k) for k in cls.__mro__).
• register_model._walk now gates the register_pytree_node block on it. Leaf types (int/float/str/bool/NoneType/tuple) and any purely __slots__ class are left as opaque JAX leaves — never registered. Real model-component classes (which have __dict__) register as before. This is a producer-side fix (per [[feedback_no_silent_guards]]) — no silent fallback inside _partition.

autofit_workspace_test (1 new file)

• scripts/jax_assertions/pytree_leaf_registration.py — reproduces af.Model(int); asserts leaf types are NOT registered, the real component IS, and tree_flatten/tree_unflatten round-trips. Lives here because JAX unit tests were moved out of test_autofit/ (test: delete jax-using unit tests (moved to autofit_workspace_test) #1247).

Ship evidence

• Repro (broken): shapelet model tree_flatten → TypeError: vars() argument must have __dict__ attribute (confirmed on main).
• Fixed: same shapelet repro → tree_flatten OK; new assertion → pytree_leaf_registration: all assertions passed.
• No regression: existing jax_assertions/pytrees.py + enable_pytrees.py → both pass.
• Helper sanity: _instances_carry_dict → int/float/str → False, normal class → True.
• API changes: none (public surface unchanged; one private helper added). Scripts changed: one new workspace_test assertion.

Detection-gap note

The new assertion is deliberately not added to autofit_workspace_test/smoke_tests.txt — the sibling jax_assertions aren't smoke-listed either, and growing curated lists is discouraged ([[feedback_smoke_tests_small_subset]]). Making the nightly actually exercise this class of JAX failure is the scope of the verify_install / harness bug (Cluster C).

On your go-ahead

Reply here (or relaunch) and I'll run the full ship_library gate (worktree pytest on PyAutoFit + review + Heart), commit, push, and open the library-first PR on PyAutoFit, then the autofit_workspace_test PR behind it.

[Jammy2211]

✅ Shipped (PRs open, pending-release). Human ship sign-off given ("ship all").

• Library: fix: don't register JAX-leaf types (int) as pytree nodes (shapelet vmap crash) #1366 — pytest test_autofit/ 1473 passed, 1 skipped; shapelet repro now flattens; existing jax_assertions no-regression
• Regression: test: JAX pytree leaf-type registration guard (guards PyAutoFit#1366) autofit_workspace_test#42 (leaf-registration guard, merge behind fix: don't register JAX-leaf types (int) as pytree nodes (shapelet vmap crash) #1366)

Heart was RED at ship (release-hold + stale-evidence, not branch breakage); shipped on explicit human go-ahead. Merge/close remain yours.