Tests

clean up #44

Job	Run time
test (ubuntu-latest, 3.10, jax)	4m 55s
lint	1m 16s
test (ubuntu-latest, 3.10, tensorflow)	4m 51s
test (ubuntu-latest, 3.10, torch)	4m 36s
test (ubuntu-latest, 3.11, jax)	4m 54s
test (ubuntu-latest, 3.11, tensorflow)	4m 54s
test (ubuntu-latest, 3.11, torch)	4m 52s
test (windows-latest, 3.10, jax)	4m 39s
test (windows-latest, 3.10, tensorflow)	1m 33s
test (windows-latest, 3.10, torch)	1m 27s
test (windows-latest, 3.11, jax)	1m 32s
test (windows-latest, 3.11, tensorflow)	1m 24s
test (windows-latest, 3.11, torch)	1m 25s
	42m 18s

Provide feedback