reproduce nonfinite md simulations when using ts.integrate

[curtischong]

Summary

This is my versioned recreation of #579

single-points are good

Single-point check:
  n outputs: 16
  energy finite: [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True]
  forces finite: [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True]
  stress finite: [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True]
  energy: [[-535.3994750976562], [-529.4705200195312], [-533.9595947265625], [-534.739013671875], [-536.77587890625], [-536.6836547851562], [-532.7955932617188], [-534.208740234375], [-536.52587890625], [-538.9987182617188], [-538.0243530273438], [-537.80322265625], [-537.6126708984375], [-534.6865234375], [-534.1015625], [-536.8688354492188]]

The final output is

MD completed
  systems: 16
  precondition wall: 79.898 s
  md wall: 356.278 s
  system-steps/s: 44.909
  peak CUDA memory: 25.094 GB
  bad indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
  bad reasons: ['temperature: system 10 contains non-finite values', 'temperature: system 11 contains non-finite values', 'positions: shape (2048, 3) is not system-indexed and contains non-finite values', 'min_pair_distance: system 10 has distance None', 'min_pair_distance: system 11 has distance None']
  final temperature K: [282.53186 294.44714 278.65186 269.4097  310.83163 310.9021  294.8846
 308.16504 344.82318 271.05295       nan       nan 298.874   316.54333
 295.01144 296.3393 ]
  final energy eV: [-558.7933  -557.4314  -555.4489  -559.05054 -558.44086 -560.5101
 -557.4262  -558.8636  -557.4632  -558.1983  -116.91453 -116.91453
 -558.4159  -560.60614 -556.5257  -557.40155]
  final min distance A: [2.23122228407305, 2.2231065372954655, 2.0732411770403822, 2.2230781210905173, 2.2186643246207405, 2.2242171978787355, 2.220441242840908, 2.1127487497556396, 2.216833290367662, 2.155452843840672, None, None, 2.238543372219303, 2.220236397340652, 2.2145372746043424, 2.1422891009012734]

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[curtischong]

whoops meant to make this PR in my personal fork for visibility