GalSim-developers / JAX-GalSim

JAX port of GalSim, for parallelized, GPU accelerated, and differentiable galaxy image simulations.
Other
25 stars 3 forks source link

vendor and modify `_wraps` from `jax._src.numpy.utils` to use for galsim #93

Closed beckermr closed 3 months ago

beckermr commented 3 months ago

@ismael-mendoza noticed that jax moved their _wraps decorator. It turns out this decorator doesn't work for galsim doc strings anyways. We should vendor their utility and modify to work with galsim doc strings.