Tests

Add soft clamping to affine transform, otherwise the loss explodes ev… #546

Job	Run time
test (ubuntu-latest, 3.10, jax)	15s
test (ubuntu-latest, 3.10, numpy)	15s
test (ubuntu-latest, 3.10, tensorflow)	15s
test (ubuntu-latest, 3.10, torch)	13s
test (ubuntu-latest, 3.11, jax)	10s
test (ubuntu-latest, 3.11, numpy)	9s
test (ubuntu-latest, 3.11, tensorflow)	13s
test (ubuntu-latest, 3.11, torch)	13s
test (windows-latest, 3.10, jax)	14s
test (windows-latest, 3.10, numpy)	15s
test (windows-latest, 3.10, tensorflow)	11s
test (windows-latest, 3.10, torch)	9s
test (windows-latest, 3.11, jax)	10s
test (windows-latest, 3.11, numpy)	10s
test (windows-latest, 3.11, tensorflow)	11s
test (windows-latest, 3.11, torch)	12s
	3m 15s

Provide feedback