* Add GPU support for matrix_solve_ls by implement the default fast path a...

* Add GPU support for matrix_solve_ls by implement the default fast path a composite Python function that only relies on matmul, cholesky and triangular_solve, which all have GPU implementations.
* Enable unit test for most linear algebra ops and their gradients on GPU.
* Temporarily disable GPU test for matrix_solve gradient until a latent bug revealed by this change is fixed.
* Add a benchmark for matrix_solve_ls.

The impact on CPU speed is generally positive except for a few regressions on tiny matrices. It is significantly faster (e.g. 50X faster for 1001x1001) for single large matrices, since most of the work is now performed in multi-threaded instead of single-threaded matmul.

Shape         Nrhs   Before              After               Speedup
(4,4)           1    6.60419464111e-05    6.103515625e-05     7.58%
(4,4)           2    6.29425048828e-05	  6.5803527832e-05   -4.55%
(4,4)           4    6.38961791992e-05	  7.29560852051e-05 -14.18%
(8,4)           1    6.19888305664e-05	  6.19888305664e-05   0.00%
(8,4)           2    6.103515625e-05	  6.38961791992e-05  -4.69%
(8,4)           4    5.98430633545e-05	  6.29425048828e-05  -5.18%
(4,8)           1    6.00814819336e-05	  6.00814819336e-05   0.00%
(4,8)           2    6.29425048828e-05	  6.19888305664e-05   1.52%
(4,8)           8    6.07967376709e-05	  6.29425048828e-05  -3.53%
(10,10)         1    6.22272491455e-05	  6.50882720947e-05  -4.60%
(10,10)         2    6.31809234619e-05	  6.48498535156e-05  -2.64%
(10,10)         10   6.19888305664e-05	  7.29560852051e-05 -17.69%
(10,8)          1    6.19888305664e-05	  8.29696655273e-05 -33.85%
(10,8)          2    7.29560852051e-05	  6.50882720947e-05  10.78%
(10,8)          8    6.19888305664e-05	  7.08103179932e-05 -14.23%
(8,10)          1    6.103515625e-05	  6.103515625e-05     0.00%
(8,10)          2    6.29425048828e-05	  6.103515625e-05     3.03%
(8,10)          10   7.41481781006e-05	  6.19888305664e-05  16.40%
(16,16)         1    7.48634338379e-05	  6.29425048828e-05  15.92%
(16,16)         2    7.5101852417e-05	  7.60555267334e-05  -1.27%
(16,16)         16   6.38961791992e-05	  7.67707824707e-05 -20.15%
(16,10)         1    8.48770141602e-05	  6.103515625e-05    28.09%
(16,10)         2    6.19888305664e-05	  6.19888305664e-05   0.00%
(16,10)         10   6.38961791992e-05	  6.29425048828e-05   1.49%
(10,16)         1    6.50882720947e-05	  6.19888305664e-05   4.76%
(10,16)         2    6.29425048828e-05	  6.103515625e-05     3.03%
(10,16)         16   7.60555267334e-05	  6.50882720947e-05  14.42%
(101,101)       1    0.000195980072021	  7.60555267334e-05  61.19%
(101,101)       2    0.000201940536499	  9.20295715332e-05  54.43%
(101,101)       101  0.000378847122192	  0.00026798248291   29.26%
(101,31)        1    8.51154327393e-05	  6.38961791992e-05  24.93%
(101,31)        2    8.20159912109e-05	  7.29560852051e-05  11.05%
(101,31)        31   9.70363616943e-05	  8.70227813721e-05  10.32%
(31,101)        1    9.58442687988e-05	  6.8187713623e-05   28.86%
(31,101)        2    0.000118017196655	  7.39097595215e-05  37.37%
(31,101)        101  0.000128984451294	  0.000111103057861  13.86%
(256,256)       1    0.00152015686035	  0.000142812728882  90.61%
(256,256)       2    0.00153708457947	  0.000188112258911  87.76%
(256,256)       256  0.00419092178345	  0.00194692611694   53.54%
(256,200)       1    0.000885009765625	  0.000108003616333  87.80%
(256,200)       2    0.000905990600586	  0.000153064727783  83.11%
(256,200)       200  0.00238180160522	  0.00109219551086   54.14%
(200,256)       1    0.000879049301147	  0.000108957290649  87.61%
(200,256)       2    0.000900030136108	  0.000159025192261  82.33%
(200,256)       256  0.00287699699402	  0.00136303901672   52.62%
(1001,1001)     1    0.0573220252991	  0.00101280212402   98.23%
(1001,1001)     2    0.0591979026794	  0.00128102302551   97.84%
(1001,1001)     1001 0.205045938492	  0.0915141105652    55.37%
(1001,501)      1    0.0129158496857	  0.000411987304688  96.81%
(1001,501)      2    0.0133857727051	  0.000498056411743  96.28%
(1001,501)      501  0.0418498516083	  0.0139379501343    66.70%
(501,1001)      1    0.012974023819	  0.000403165817261  96.89%
(501,1001)      2    0.0131361484528	  0.000482082366943  96.33%
(501,1001)      1001 0.0733780860901	  0.0280270576477    61.80%
(1024,1024)     1    0.0646319389343	  0.00120306015015   98.14%
(1024,1024)     2    0.0661849975586	  0.00151205062866   97.72%
(1024,1024)     1024 0.227252960205	  0.0971050262451    57.27%
(1024,128)      1    0.00104093551636	  9.60826873779e-05  90.77%
(1024,128)      2    0.00108408927917	  0.000156879425049  85.53%
(1024,128)      128  0.00251793861389	  0.000633955001831  74.82%
(128,1024)      1    0.00103306770325	  9.41753387451e-05  90.88%
(128,1024)      2    0.00111484527588	  0.000247955322266  77.76%
(128,1024)      1024 0.0141911506653	  0.00415301322937   70.74%
(2048,2048)     1    0.47181892395	  0.00541090965271   98.85%
(2048,2048)     2    0.477843999863	  0.00651216506958   98.64%
(2048,2048)     2048 1.75637602806	  0.759104013443     56.78%
(2048,64)       1    0.000587940216064	  9.08374786377e-05  84.55%
(2048,64)       2    0.000652074813843	  0.000128030776978  80.37%
(2048,64)       64   0.00130105018616	  0.000342130661011  73.70%
(64,2048)       1    0.000602006912231	  8.60691070557e-05  85.70%
(64,2048)       2    0.000680923461914	  0.000250816345215  63.17%
(64,2048)       2048 0.027850151062	  0.00638389587402   77.08%
(513,4,4)      1    0.000416040420532	  0.000288963317871  30.54%
(513,4,4)      2    0.000466823577881	  0.000313997268677  32.74%
(513,4,4)      4    0.000488996505737	  0.000338077545166  30.86%
(513,4,2)      1    0.000366926193237	  0.0002760887146    24.76%
(513,4,2)      2    0.000352144241333	  0.000269889831543  23.36%
(513,4,2)      2    0.000343084335327	  0.000263929367065  23.07%
(513,2,4)      1    0.000388145446777	  0.000275850296021  28.93%
(513,2,4)      2    0.000396013259888	  0.000270128250122  31.79%
(513,2,4)      4    0.00041389465332	  0.000282049179077  31.85%
(513,16,16)    1    0.000340938568115	  0.000346899032593  -1.75%
(513,16,16)    2    0.0003821849823	  0.000375986099243   1.62%
(513,16,16)    16   0.000468015670776	  0.000473022460938  -1.07%
(513,16,10)    1    0.000268936157227	  0.000272989273071  -1.51%
(513,16,10)    2    0.000308990478516	  0.000308036804199   0.31%
(513,16,10)    10   0.000367164611816	  0.000363111495972   1.10%
(513,10,16)    1    0.0002601146698	  0.000262975692749  -1.10%
(513,10,16)    2    0.000304937362671	  0.000303983688354   0.31%
(513,10,16)    16   0.00036883354187	  0.000371932983398  -0.84%
(513,256,256)  1    0.109347105026	  0.11101102829	     -1.52%
(513,256,256)  2    0.114288806915	  0.114917993546     -0.55%
(513,256,256)  256  0.353310108185	  0.348378896713      1.40%
(513,256,128)  1    0.0285301208496	  0.0287981033325    -0.94%
(513,256,128)  2    0.0313191413879	  0.0301601886749     3.70%
(513,256,128)  128  0.0773530006409	  0.0720520019531     6.85%
(513,128,256)  1    0.0281269550323	  0.0283789634705    -0.90%
(513,128,256)  2    0.0297160148621	  0.0298500061035    -0.45%
(513,128,256)  256  0.140498876572	  0.126143932343     10.22%

PiperOrigin-RevId: 167026381