jax

3d6b521d - The reshape call has a numerical bug where the target reshape size is always higher than array size if num_devices is not a nonnegative power of two. To fix this, this commit adds padding