gh-903: Port galaxies.py to support jax

[connoraird]

Description

Ports functions in galaxies.py which mutate arrays to now use xpx.at... so they support jax.

Closes: #903

Changelog entry

Added: jax support for glass.galaxies.redshifts_from_nz and glass.galaxies.galaxy_shear .

Checks

• Is your code passing linting?
• Is your code passing tests?
• Have you added additional tests (if required)?
• Have you modified/extended the documentation (if required)?
• Have you added a one-liner changelog entry above (if required)?

[connoraird]

1. Coveralls reduction I assume is because of the removal of array-api-strict
2. I believe this could be a genuine regression in performance

[paddyroddy]

Looks good. How do we know it's a genuine reduction this time?

[connoraird]

Looks good. How do we know it's a genuine reduction this time?

I'm assuming it is, as the only tests which seem to regularly show up as regressed is test_galaxy_shear which is exactly what I've updated.

[paddyroddy]

That would make sense. Thankfully not a huge regression.

Performance has regressed:
	test_galaxy_shear[numpy-True] (0001_f906eb0) - Field 'mean' has failed PercentageRegressionCheck: 9.996939723 > 5.000000000
	test_galaxy_shear[numpy-False] (0001_f906eb0) - Field 'mean' has failed PercentageRegressionCheck: 8.899524405 > 5.000000000