Skip to content

fix(jax): defensive pytree dedup in imaging/interferometer analyses#561

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/release-unblockers
May 29, 2026
Merged

fix(jax): defensive pytree dedup in imaging/interferometer analyses#561
Jammy2211 merged 1 commit into
mainfrom
feature/release-unblockers

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Same defensive pattern as the PyAutoGalaxy companion PR — additionally cross-populates `Tracer` from autofit's `_REGISTERED_INSTANCE_CLASSES` since lens models always carry one, making `Tracer` a regular source of duplicate-registration whenever both the autofit walker and the analysis registration path run.

Repro

Same root cause as autogalaxy companion (multi/start_here.py in autogalaxy_workspace during the 2026.5.29.2 release):

```
ValueError: Duplicate custom PyTreeDef type registration for
<class 'autoarray.dataset.dataset_model.DatasetModel'>.
```

The autolens analyses register the same class set plus `Tracer`. Once `autofit.jax.pytrees.register_model` walks a lens model, both `DatasetModel` and `Tracer` get registered via autofit's path — neither set lands in autoarray's `_pytree_registered_classes`, so the next `fit_from` call's registration step asks JAX to register them again.

Fix

For both `imaging/model/analysis.py` and `interferometer/model/analysis.py`:

  1. Module-level `FIT*_PYTREES_REGISTERED = False` flag.
  2. Cross-populate `(DatasetModel, Tracer)` from autofit's set into autoarray's.
  3. Set the flag at the end.

Paired with

  • PyAutoGalaxy companion PR (same fix)
  • PyAutoBuild PR (release.yml `run_scripts` env block)

Test plan

  • `pytest test_autolens/imaging/model test_autolens/interferometer/model` — passed (unit tests don't exercise JAX path)
  • CI on this PR
  • Next `autobuild pre_build` dispatch

🤖 Generated with Claude Code

Same defense as PyAutoGalaxy companion PR — but additionally
cross-populates ``Tracer`` from autofit's
``_REGISTERED_INSTANCE_CLASSES`` since lens models always carry one,
making ``Tracer`` a regular source of duplicate-registration when both
the autofit walker and the analysis registration path run.

Triggered by ``multi/start_here.py`` in autogalaxy_workspace during the
2026.5.29.2 release; same `ValueError: Duplicate custom PyTreeDef type
registration for DatasetModel` would surface here as soon as a
multi-dataset lens model gets fit.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@Jammy2211 Jammy2211 merged commit edca231 into main May 29, 2026
6 checks passed
@Jammy2211 Jammy2211 deleted the feature/release-unblockers branch May 29, 2026 09:39
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