Weak-lensing visualization crashes on JAX path: fit.shear_yx lacks .grid #614

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.