Wassim KABALAN 6693e5c725 Fix sharding error (#37) ... * Use cosmo as arg for the ODE function * Update examples * format * notebook update * fix tests * add correct annotations for weights in painting and warning for cic_paint in distributed pm * update test_against_fpm * update distributed tests and add jacfwd jacrev and vmap tests * format * add Caveats to notebook readme * final touches * update Growth.py to allow using FastPM solver * fix 2D painting when input is (X , Y , 2) shape * update cic read halo size and notebooks examples * Allow env variable control of caching in growth * Format * update test jax version * update notebooks/03-MultiGPU_PM_Halo.ipynb * update numpy install in wf * update tolerance :) * reorganize install in test workflow * update tests * add mpi4py * update tests.yml * update tests * update wf * format * make normal_field signature consistent with jax.random.normal * update by default normal_field dtype to match JAX * format * debug test workflow * format * debug test workflow * updating tests * fix accuracy * fixed tolerance * adding caching * Update conftest.py * Update tolerance and precision settings in distributed PM tests * revererting back changes to growth.py --------- Co-authored-by: Francois Lanusse <fr.eiffel@gmail.com> Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>		2025-06-28 23:07:31 +02:00
..
conftest.py	Fix sharding error (#37)	2025-06-28 23:07:31 +02:00
helpers.py	jaxdecomp proto (#21)	2024-12-20 05:44:02 -05:00
test_against_fpm.py	Fix sharding error (#37)	2025-06-28 23:07:31 +02:00
test_distributed_pm.py	Fix sharding error (#37)	2025-06-28 23:07:31 +02:00
test_gradients.py	Fix sharding error (#37)	2025-06-28 23:07:31 +02:00