ENH update to latest JAX version and slight refactor in test suite - Githubissues

GalSim-developers / JAX-GalSim

JAX port of GalSim, for parallelized, GPU accelerated, and differentiable galaxy image simulations.

Other

25 stars 3 forks source link

ENH update to latest JAX version and slight refactor in test suite #94

Closed ismael-mendoza closed 3 months ago

ismael-mendoza commented 3 months ago

I was thinking it would be good to update to the latest jax version and pin jax-galsim to that if possible. Open to discussion on this though.

Changes:

_wraps to implements (still won't work for galsim but allows code to run, see #93 )
pin tensorflow-probability, as the newest version does not have the substrates packages?

beckermr commented 3 months ago

tfp changed their default submodule imports. A direct import of the bessel function we need works fine.

beckermr commented 3 months ago

Next issue here is that pytest version 8 broke the test mangling we do on galsim. I am looking into that now.

beckermr commented 3 months ago

This PR is ready to go. No API changes here, just updates to the tests etc. Shall we merge @ismael-mendoza?

beckermr commented 3 months ago

Also, we should squash when merging. The commit history here is a disaster.

ismael-mendoza commented 3 months ago

yes I think so! thank you!

I agree with the squash and merge, I removed the other options from the settings

beckermr commented 3 months ago

Ahhhhh we should keep the other options in the settings IMHO. I simply meant for this PR we'd squash.