fix imports #632

Job	Run time
test (ubuntu-latest, 3.10, jax)	1m 56s
test (ubuntu-latest, 3.10, tensorflow)	2m 45s
test (ubuntu-latest, 3.10, torch)	2m 20s
test (ubuntu-latest, 3.11, jax)	2m 3s
test (ubuntu-latest, 3.11, tensorflow)	2m 44s
test (ubuntu-latest, 3.11, torch)	2m 29s
test (windows-latest, 3.10, jax)	7m 38s
test (windows-latest, 3.10, tensorflow)	7m 37s
test (windows-latest, 3.10, torch)	8m 7s
test (windows-latest, 3.11, jax)	5m 57s
test (windows-latest, 3.11, tensorflow)	8m 49s
test (windows-latest, 3.11, torch)	10m 13s
	1h 2m 38s

Provide feedback