feat(jax): force & virial (#4251) #3227

Job	Run time
Test Python (1, 3.9)	26m 27s
Test Python (1, 3.12)	31m 38s
Test Python (2, 3.9)	26m 41s
Test Python (2, 3.12)	31m 45s
Test Python (3, 3.9)	28m 56s
Test Python (3, 3.12)	33m 39s
Test Python (4, 3.9)	28m 6s
Test Python (4, 3.12)	32m 32s
Test Python (5, 3.9)	28m 5s
Test Python (5, 3.12)	31m 18s
Test Python (6, 3.9)	28m 21s
Test Python (6, 3.12)	33m 35s
Combine and update integration test durations (3.9)	31s
Combine and update integration test durations (3.12)	55s
Pass testing Python	4s
	6h 2m 33s

Provide feedback