Bugfix concatenation axis #629

Job	Run time
test (ubuntu-latest, 3.10, jax)	2m 3s
test (ubuntu-latest, 3.10, tensorflow)	2m 33s
test (ubuntu-latest, 3.10, torch)	2m 29s
test (ubuntu-latest, 3.11, jax)	2m 1s
test (ubuntu-latest, 3.11, tensorflow)	2m 40s
test (ubuntu-latest, 3.11, torch)	2m 28s
test (windows-latest, 3.10, jax)	7m 18s
test (windows-latest, 3.10, tensorflow)	8m 42s
test (windows-latest, 3.10, torch)	6m 39s
test (windows-latest, 3.11, jax)	5m 52s
test (windows-latest, 3.11, tensorflow)	7m 39s
test (windows-latest, 3.11, torch)	6m 57s
	57m 21s

Provide feedback