Tests

Cherry picked simulator test #16

Job	Run time
lint	1m 18s
test (ubuntu-latest, 3.10, jax)	2m 13s
test (ubuntu-latest, 3.10, tensorflow)	3m 22s
test (ubuntu-latest, 3.10, torch)	3m 21s
test (ubuntu-latest, 3.11, jax)	3m 3s
test (ubuntu-latest, 3.11, tensorflow)	3m 18s
test (ubuntu-latest, 3.11, torch)	3m 20s
test (windows-latest, 3.10, jax)	3m 6s
test (windows-latest, 3.10, tensorflow)	3m 5s
test (windows-latest, 3.10, torch)	3m 4s
test (windows-latest, 3.11, jax)	3m 4s
test (windows-latest, 3.11, tensorflow)	3m 3s
test (windows-latest, 3.11, torch)	3m 5s
	38m 22s

Provide feedback