guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-29 00:21:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Small fix

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-22 13:00:25 -04:00
parent 85cca44fb0
commit d28982eec7
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
jaxpm/pm.py
View file

@ -157,7 +157,8 @@ def make_ode_fn(mesh_shape, halo_size=0, sharding=None):
return nbody_ode return nbody_ode
def get_ode_fn(cosmo:Cosmology, mesh_shape):
def get_ode_fn(cosmo, mesh_shape):
def nbody_ode(a, state, args): def nbody_ode(a, state, args):
""" """