PR #10388: Address lost changes for pr 7854
PR #10388: Address lost changes for pr 7854

Imported from GitHub PR https://github.com/openxla/xla/pull/10388


Copybara import of the project:

--
d55af60f1335d82fc9c32d0dfb34708cb32174a3 by TJ <tjx@nvidia.com>:

Address lost changes for pr https://github.com/openxla/xla/pull/7854

Merging this change closes #10388

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10388 from Tixxx:tixxx/coll_matmul_fix d55af60f1335d82fc9c32d0dfb34708cb32174a3
