JAX Aggregator issue

[Jammy2211]

autoconf -> main
autofit -> feature/jax_mode
autoarray -> feature/jax_wrapper
autogalaxy -> feature/jax_wrapper

Thiss group of repos with the JAX version (e.g. USE_JAX=1 ) have the following error for the autogalaxy/aggregator tests:

test_autogalaxy/aggregator/interferometer/test_aggregator_fit_interferometer.py:135:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../PyAutoFit/autofit/database/aggregator/aggregator.py:255: in map
    yield func(fit)
../PyAutoFit/autofit/aggregator/base.py:162: in func_gen
    self.object_via_gen_from(
autogalaxy/aggregator/interferometer/fit_interferometer.py:159: in object_via_gen_from
    return _fit_interferometer_from(
autogalaxy/aggregator/interferometer/fit_interferometer.py:65: in _fit_interferometer_from
    adapt_images_list = agg_util.adapt_images_from(fit=fit)
autogalaxy/aggregator/agg_util.py:122: in adapt_images_from
    if fit.value(name="adapt_images") is None:
../PyAutoFit/autofit/database/model/fit.py:505: in value
    return self.__getitem__(item=name)
../PyAutoFit/autofit/database/model/fit.py:364: in __getitem__
    value = p.value
../PyAutoFit/autofit/database/model/array.py:116: in value
    return self.hdu_list
../PyAutoFit/autofit/database/model/array.py:108: in hdu_list
    return fits.HDUList([hdu.hdu for hdu in self.hdus])
../PyAutoFit/autofit/database/model/array.py:179: in hdu
    self.array,
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <HDU None>

    @property
    def array(self):
>       return np.frombuffer(
            self.bytes,
            dtype=self.dtype,
        ).reshape(self.shape)
E       ValueError: buffer size must be a multiple of element size

../PyAutoFit/autofit/database/model/array.py:49: ValueError

The strange thing is these tests don't fail on autolens .