pmap_test::PythonPmapTest::testRule30

[rahulbatra85]

https://github.com/google/jax/blob/main/tests/pmap_test.py#L1128

Isn't hardcoding the length of array to 40 an issue if num_devices doesn't evenly divide 40. The reshape operation will fail. Right?

I am running on a machine with 16 AMD GPUs and reshape operation fails. However, if I reduce the available GPUs to 8 then it works. Trying to understand if this is an issue with the testcase or somewhere else?

[jakevdp]

Hi - thanks for the question. This particular set of tests is designed to be run with 8 devices: https://github.com/google/jax/blob/9a0de29114e02657a552210c5a080d49ec16a88b/tests/pmap_test.py#L86-L96

If you run the test with other device counts, then I'd expect you'll get some failures.

[rahulbatra85]

ah ok, I missed that!

[hawkinsp]

That said, I think it's fine to improve the tests so they work with different numbers.

e.g., here dividing into gcd(num devices you have, 40) would probably work.