Bug
Running AnalysisWeak(dataset, use_jax=True) with a gradient search (e.g. af.MultiStartAdam) crashes during the search's visualization update:
File "autolens/weak/plot/weak_dataset_plots.py", line 34, in _positions_yx
grid = shear_yx.grid
AttributeError: 'jaxlib._jax.ArrayImpl' object has no attribute 'grid'
On the JAX path the max-likelihood FitWeak's shear_yx reaching the plotter is a raw JAX array rather than the structured object carrying .grid, so subplot_fit_weak → plot_shear_yx_2d → _positions_yx fails.
Reproduced on the RAL A100 (jax 0.10.2, current mains) via autolens_jax_joss/benchmarks/weak.py (A2744 pyRRG catalogue, NFWSph model) — full traceback in the job log referenced from PyAutoLabs/autolens_workspace#281. The NumPy path (use_jax=False, as scripts/weak/real_data/a2744.py uses with Nautilus) is unaffected.
Workaround: PYAUTO_SKIP_VISUALIZATION=1 (the benchmark harness now sets this for timing purity regardless).
Likely fix: the weak visualizer should rebuild/attach the positions grid when the fit comes off the JAX path (or _positions_yx should fall back to the dataset's grid) before plotting.
Bug
Running
AnalysisWeak(dataset, use_jax=True)with a gradient search (e.g.af.MultiStartAdam) crashes during the search's visualization update:On the JAX path the max-likelihood
FitWeak'sshear_yxreaching the plotter is a raw JAX array rather than the structured object carrying.grid, sosubplot_fit_weak→plot_shear_yx_2d→_positions_yxfails.Reproduced on the RAL A100 (jax 0.10.2, current mains) via
autolens_jax_joss/benchmarks/weak.py(A2744 pyRRG catalogue, NFWSph model) — full traceback in the job log referenced from PyAutoLabs/autolens_workspace#281. The NumPy path (use_jax=False, asscripts/weak/real_data/a2744.pyuses with Nautilus) is unaffected.Workaround:
PYAUTO_SKIP_VISUALIZATION=1(the benchmark harness now sets this for timing purity regardless).Likely fix: the weak visualizer should rebuild/attach the positions grid when the fit comes off the JAX path (or
_positions_yxshould fall back to the dataset's grid) before plotting.