Replace Python for loop with jax.lax.scan in step_n

[maximilianGord]

Currently, step_n uses a Python for loop to iterate over Langevin steps. While correct, this causes JAX to unroll the full loop at trace/compile time, which leads to longer compilation times and higher memory usage as num_steps grows.

This PR replaces the loop with jax.lax.scan, which compiles the loop body once regardless of num_steps, making compilation time constant and reduces memory.

The behavior is identical — same random key splitting logic, same step_single calls, same outputs. The **kwargs are captured via closure inside scan_step, which works correctly as long as kwargs don't vary per step (consistent with the existing usage).

[plainerman]

Dear @maximilianGord
Thank you for the PR!

We already do jax.lax.scan for the simulation see here:

ScoreMD/src/scoremd/simulation.py

Line 68 in 149b1ec

_, (trajectory, velocities) = jax.lax.scan(step_fn, init_carry, None, length=n_steps)

The problem with jax.lax.scan is that it launches a new cuda kernel for every run, meaning that it reduces overhead. In my tests it was faster if we compile a few steps by unrolling num_steps (like 10-50 steps) and then do the actual simulation by producing this in batch.

This also reduces memory allocations since the arrays can actually be overridden.

Have you performed any benchmarks?

As of now, I do not see what this change would add to the code base. If you think compilation time is actually a problem, we can talk about just setting num_steps to 1 (or a smaller number) in the default configs and then handling everything via n_steps.