jax

85f30913 - The reshape call has a numerical bug where the target reshape size is always higher than array size if num_devices is not a multiple of two. To fix this, this commit skips the test if the condition doesn't hold