guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-23 01:57:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

format

Browse source
This commit is contained in:
Wassim Kabalan 2024-12-21 23:27:05 +01:00
parent a924458f0d
commit bbacd45dcf
1 changed files with 6 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
tests/test_gradients.py
View file

@ -1,6 +1,7 @@
import jax
import pytest
from diffrax import Dopri5, ODETerm, PIDController, SaveAt, diffeqsolve , RecursiveCheckpointAdjoint, BacksolveAdjoint
from diffrax import (BacksolveAdjoint, Dopri5, ODETerm, PIDController,
RecursiveCheckpointAdjoint, SaveAt, diffeqsolve)
from helpers import MSE
from jax import numpy as jnp
@ -23,7 +24,8 @@ def test_nbody_grad(simulation_config, initial_conditions, lpt_scale_factor,
if adjoint == 'OTD':
pytest.skip("OTD adjoint not implemented yet (needs PFFT3D JVP)")
adjoint = RecursiveCheckpointAdjoint() if adjoint == 'DTO' else BacksolveAdjoint(solver=Dopri5())
adjoint = RecursiveCheckpointAdjoint(
) if adjoint == 'DTO' else BacksolveAdjoint(solver=Dopri5())
@jax.jit
@jax.grad