Tests

allow int as type for batch_shape #224

Job	Run time
lint	1m 17s
test (ubuntu-latest, 3.10, jax)	5m 38s
test (ubuntu-latest, 3.10, tensorflow)	5m 29s
test (ubuntu-latest, 3.10, torch)	4m 14s
test (ubuntu-latest, 3.11, jax)	6m 39s
test (ubuntu-latest, 3.11, tensorflow)	5m 43s
test (ubuntu-latest, 3.11, torch)	4m 24s
test (windows-latest, 3.10, jax)	11m 19s
test (windows-latest, 3.10, tensorflow)	11m 11s
test (windows-latest, 3.10, torch)	10m 27s
test (windows-latest, 3.11, jax)	12m 3s
test (windows-latest, 3.11, tensorflow)	13m 28s
test (windows-latest, 3.11, torch)	10m 24s
	1h 42m 16s

Provide feedback