Tests

Update dependencies #390

Job	Run time
lint	1m 30s
test (ubuntu-latest, 3.10, jax)	7m 44s
test (ubuntu-latest, 3.10, tensorflow)	8m 8s
test (ubuntu-latest, 3.10, torch)	5m 11s
test (ubuntu-latest, 3.11, jax)	8m 26s
test (ubuntu-latest, 3.11, tensorflow)	9m 47s
test (ubuntu-latest, 3.11, torch)	5m 34s
test (windows-latest, 3.10, jax)	13m 52s
test (windows-latest, 3.10, tensorflow)	14m 18s
test (windows-latest, 3.10, torch)	13m 43s
test (windows-latest, 3.11, jax)	14m 40s
test (windows-latest, 3.11, tensorflow)	13m 42s
test (windows-latest, 3.11, torch)	12m 32s
	2h 9m 7s

Provide feedback