Tests

Enable squeeze option by broadcast #416

Job	Run time
lint	1m 22s
test (ubuntu-latest, 3.10, jax)	8m 10s
test (ubuntu-latest, 3.10, tensorflow)	9m 3s
test (ubuntu-latest, 3.10, torch)	5m 47s
test (ubuntu-latest, 3.11, jax)	8m 51s
test (ubuntu-latest, 3.11, tensorflow)	9m 47s
test (ubuntu-latest, 3.11, torch)	6m 25s
test (windows-latest, 3.10, jax)	14m 58s
test (windows-latest, 3.10, tensorflow)	15m 35s
test (windows-latest, 3.10, torch)	12m 32s
test (windows-latest, 3.11, jax)	14m 52s
test (windows-latest, 3.11, tensorflow)	17m 42s
test (windows-latest, 3.11, torch)	13m 13s
	2h 18m 17s

Provide feedback