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
Loading
Please sign in to comment