Commit f59a82f2 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower Committed by TensorFlower Gardener
Browse files

Replacing the current inner Cholesky decomposition loop with a While loop rolled version.

This will allow for much larger Cholesky decompositions (and thus matrix inversions) than previously possible on TPU because of the use of rolled While loops so XLA compilation will no longer timeout.

While there is a minor runtime performance decrease (now 25ms vs 15ms for a 500x500 matrix) the compilation time is significantly faster (12.8s vs 55.2s for a 500x500 matrix.)

PiperOrigin-RevId: 193114816
parent 39635af4
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment