Fixed internal issue.
Fixed internal issue.

This is follow up to https://github.com/openxla/xla/commit/0ab2be0b5a575da3206d2c2f92b85e6346708405, 
which got reverted here https://github.com/openxla/xla/commit/b4a46471fd1b80452f913d927e7db04d4a6e260b

Reverts 5c66d026b089c47b021a65663153a268f6598e03
