0a7be1732c
This increases performance, particularly for the Newton solver on TPU. PiperOrigin-RevId: 600696483 Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff