Tests

Fix mlp / hidden block #499

Job	Run time
lint	1m 23s
test (ubuntu-latest, 3.10, jax)	7m 22s
test (ubuntu-latest, 3.10, tensorflow)	17m 25s
test (ubuntu-latest, 3.10, torch)	8m 23s
test (ubuntu-latest, 3.11, jax)	8m 7s
test (ubuntu-latest, 3.11, tensorflow)	18m 44s
test (ubuntu-latest, 3.11, torch)	8m 51s
test (windows-latest, 3.10, jax)	13m 14s
test (windows-latest, 3.10, tensorflow)	26m 29s
test (windows-latest, 3.10, torch)	17m 8s
test (windows-latest, 3.11, jax)	11m 38s
test (windows-latest, 3.11, tensorflow)	1m 36s
test (windows-latest, 3.11, torch)	15m 3s
	2h 35m 23s

Provide feedback