Merge pull request #8 from abhijeetgangan/revert-5-Tersoff_unit_test

Revert "Try #1"
abhijeetgangan · May 22, 2024 · 90065ff · 90065ff
2 parents 0e26bea + 395aac1
commit 90065ff
Showing 1 changed file with 0 additions and 8 deletions.
diff --git a/tests/rigid_body_test.py b/tests/rigid_body_test.py
@@ -249,8 +249,6 @@ def total_energy(state, nbrs):
     nbrs = neighbor_fn.allocate(body)
     state = init_fn(key, body, 1e-3, mass=shape.mass(), neighbor=nbrs)
     E_initial = total_energy(state, nbrs)
-
-    @jit
     def step(i, state_nbrs):
       state, nbrs = state_nbrs
       nbrs = nbrs.update(state.position)
@@ -303,8 +301,6 @@ def total_energy(state, nbrs):
     nbrs = neighbor_fn.allocate(body)
     state = init_fn(key, body, 1e-3, mass=shape.mass(), neighbor=nbrs)
     E_initial = total_energy(state, nbrs)
-
-    @jit
     def step(i, state_nbrs):
       state, nbrs = state_nbrs
       nbrs = nbrs.update(state.position)
@@ -376,8 +372,6 @@ def total_energy(state, nbrs):
                     mass=shape.mass(shape_species),
                     neighbor=nbrs)
     E_initial = total_energy(state, nbrs)
-
-    @jit
     def step(i, state_nbrs):
       state, nbrs = state_nbrs
       nbrs = nbrs.update(state.position)
@@ -860,7 +854,6 @@ def test_nvt_2d_neighbor_list(self, dtype):
     state = init_fn(key, body, mass=shape.mass(), neighbor=nbrs)
     E_initial = simulate.nvt_nose_hoover_invariant(energy_fn, state, kT,
                                                    neighbor=nbrs)
-    @jit
     def sim_fn(i, state_nbrs):
       state, nbrs = state_nbrs
       state = step_fn(state, neighbor=nbrs)
@@ -918,7 +911,6 @@ def test_nvt_2d_multi_atom_species_neighbor_list(self, dtype):
     state = init_fn(key, body, mass=shape.mass(), neighbor=nbrs)
     E_initial = simulate.nvt_nose_hoover_invariant(energy_fn, state, kT,
                                                    neighbor=nbrs)
-    @jit
     def sim_fn(i, state_nbrs):
       state, nbrs = state_nbrs
       state = step_fn(state, neighbor=nbrs)