Tests

Merge pull request #263 from eodole/dev #374

Job	Run time
lint	1m 15s
test (ubuntu-latest, 3.10, jax)	8m 0s
test (ubuntu-latest, 3.10, tensorflow)	7m 59s
test (ubuntu-latest, 3.10, torch)	5m 20s
test (ubuntu-latest, 3.11, jax)	10m 8s
test (ubuntu-latest, 3.11, tensorflow)	8m 28s
test (ubuntu-latest, 3.11, torch)	5m 20s
test (windows-latest, 3.10, jax)	14m 32s
test (windows-latest, 3.10, tensorflow)	14m 19s
test (windows-latest, 3.10, torch)	13m 54s
test (windows-latest, 3.11, jax)	15m 53s
test (windows-latest, 3.11, tensorflow)	13m 54s
test (windows-latest, 3.11, torch)	10m 51s
	2h 9m 53s

Provide feedback