Standardize parameter typing to float | torch.Tensor when appropriate

[CompRhys]

see #325

I am unsure if step functions should allow floats because then we have overhead at the integration level. Init functions are a one off cost not a linear cost.

[CompRhys]

the changes here make the code less clean.

[orionarcher]

LGTM

[CompRhys]

Will wait on a speed check from @falletta and then merge if it looks fine. I can't imagine it's significant but no real harm in checking.

[janosh]

@CompRhys in case this gets merged today, could you apply

diff --git a/torch_sim/models/fairchem.py b/torch_sim/models/fairchem.py
--- a/torch_sim/models/fairchem.py
+++ b/torch_sim/models/fairchem.py
@@ -201,11 +201,11 @@ class FairChemModel(ModelInterface):
             zip(n_atoms, torch.cumsum(n_atoms, dim=0), strict=False)
         ):
             # Extract system data
-            positions = sim_state.positions[c - n : c].cpu().numpy()
-            atomic_nums = sim_state.atomic_numbers[c - n : c].cpu().numpy()
-            pbc = sim_state.pbc.cpu().numpy()
+            positions = sim_state.positions[c - n : c].detach().cpu().numpy()
+            atomic_nums = sim_state.atomic_numbers[c - n : c].detach().cpu().numpy()
+            pbc = sim_state.pbc.detach().cpu().numpy()
             cell = (
-                sim_state.row_vector_cell[idx].cpu().numpy()
+                sim_state.row_vector_cell[idx].detach().cpu().numpy()
                 if sim_state.row_vector_cell is not None
                 else None
             )

to close #469. could add a single-step fairchem relax test but if models get moved to external repos, maybe not worth it

[CompRhys]

to close #469. could add a single-step fairchem relax test but if models get moved to external repos, maybe not worth it

The fairchem PR I made to upstream has been sitting for a while without eyes from meta so I will add here.

[falletta]

I see similar performance when running the 8_benchmarking.py script. Below are the results, rounded to the first digit, obtained both before and after the changes:

=== Static benchmark, n=5000 static_time=7.9 s
=== Relax benchmark, n=500 relax_10_time=34.2 s
=== NVE benchmark, n=500 nve_time=4.5 s
=== NVT benchmark, n=500 nvt_time=4.7 s

I also checked the scalings of various operations with system size using my own profiling scripts, and everything looks good. So green light from my side @CompRhys