-
Pedro Cuenca authored
Flax tests: don't hardcode number of devices. This makes it possible to test on CPU/GPU. However, expected slices are only checked when there are 8 devices.
af279434
Flax tests: don't hardcode number of devices. This makes it possible to test on CPU/GPU. However, expected slices are only checked when there are 8 devices.