Add tf.Cholesky and tf.MatrixTriangularSolve to Flex

closes tensorflow#46577

PiperOrigin-RevId: 353135213
Change-Id: I27d78947b070d2ec64820f473466870e71745bfc

Author: abattery

tensorflow/core/kernels/BUILD

@@ -5974,6 +5974,7 @@ filegroup(
         "//tensorflow/core/kernels/linalg:linalg_ops_common.h",
         "//tensorflow/core/kernels/linalg:matrix_diag_op.h",
         "//tensorflow/core/kernels/linalg:matrix_set_diag_op.h",
+        "//tensorflow/core/kernels/linalg:matrix_triangular_solve_op_impl.h",
         "//tensorflow/core/util:image_resizer_state.h",
     ],
 )
@@ -6280,9 +6281,12 @@ filegroup(
         "//tensorflow/core/kernels/image:resize_nearest_neighbor_op.cc",
         "//tensorflow/core/kernels/image:sample_distorted_bounding_box_op.cc",
         "//tensorflow/core/kernels/linalg:linalg_ops_common.cc",
+        "//tensorflow/core/kernels/linalg:cholesky_op.cc",
         "//tensorflow/core/kernels/linalg:matrix_diag_op.cc",
         "//tensorflow/core/kernels/linalg:matrix_inverse_op.cc",
         "//tensorflow/core/kernels/linalg:matrix_set_diag_op.cc",
+        "//tensorflow/core/kernels/linalg:matrix_triangular_solve_op_complex.cc",
+        "//tensorflow/core/kernels/linalg:matrix_triangular_solve_op_real.cc",
     ],
 )
 

tensorflow/core/kernels/linalg/BUILD

@@ -34,6 +34,7 @@ package_group(
 
 # Export a few files for use on Android.
 exports_files([
+    "cholesky_op.cc",
     "einsum_op_impl_half.cc",
     "einsum_op_impl_bfloat16.cc",
     "einsum_op_impl_int32.cc",
@@ -51,6 +52,9 @@ exports_files([
     "matrix_inverse_op.cc",
     "matrix_set_diag_op.h",
     "matrix_set_diag_op.cc",
+    "matrix_triangular_solve_op_complex.cc",
+    "matrix_triangular_solve_op_impl.h",
+    "matrix_triangular_solve_op_real.cc",
 ])
 
 # Public support libraries ----------------------------------------------------

tensorflow/lite/delegates/flex/allowlisted_flex_ops.cc

@@ -67,12 +67,14 @@ const std::set<std::string>& GetFlexAllowlist() {
           "AvgPool3D",
           "AvgPool3DGrad",
           "AvgPoolGrad",
+          "BatchCholesky",
           "BatchMatMul",
           "BatchMatMulV2",
           "BatchMatrixDiag",
           "BatchMatrixDiagPart",
           "BatchMatrixInverse",
           "BatchMatrixSetDiag",
+          "BatchMatrixTriangularSolve",
           "BatchNormWithGlobalNormalization",
           "BatchNormWithGlobalNormalizationGrad",
           "BatchToSpace",
@@ -104,6 +106,7 @@ const std::set<std::string>& GetFlexAllowlist() {
           "Ceil",
           "CheckNumerics",
           "CheckNumericsV2",
+          "Cholesky",
           "CombinedNonMaxSuppression",
           "Complex",
           "ComplexAbs",
@@ -264,6 +267,7 @@ const std::set<std::string>& GetFlexAllowlist() {
           "MatrixSetDiag",
           "MatrixSetDiagV2",
           "MatrixSetDiagV3",
+          "MatrixTriangularSolve",
           "Max",
           "MaxPool",
           "MaxPool3D",