[Feature Request] MatrixSquareRoot for tflite

[gaikwadrahul8]

Tensorflow lite currently will not compile graphs containing tf.linalg.sqrtm. If you try, you get

Some ops in the model are custom ops, See instructions to implement custom ops: https://www.tensorflow.org/lite/guide/ops_custom 
Custom ops: MatrixSquareRoot
Details:
	tf.MatrixSquareRoot(tensor<3x3xf32>) -> (tensor<3x3xf32>) : {T = f32, device = ""}

As a work around, I am using an approximate solution:

def tf_denmann_beavers_sqrtm(matrix: TensorInvCovMat, n_iter=10):
    """
    Approximate the matrix-square-root by Denmann Beavers iteration
        https://en.wikipedia.org/wiki/Square_root_of_a_matrix#By_Denman%E2%80%93Beavers_iteration
    Convergence is not guaranteed.  Use at your own risk!
    """
    ym = matrix
    zm = tf.eye(tf.shape(matrix[0])[0], dtype=matrix.dtype)
    for i in range(n_iter):
        ym_ = 0.5 * (ym + tf.linalg.inv(zm))
        zm = 0.5 * (zm + tf.linalg.inv(ym))
        ym = ym_
    return ym

But it is not ideal.

Can we please get tf.linalg.sqrtm included as a standard tflite op?

[gaikwadrahul8]

This issue originally reported by @petered has been moved to this dedicated repository for LiteRT to enhance issue tracking and prioritization. To ensure continuity, we have created this new issue on your behalf.

We appreciate your understanding and look forward to your continued involvement.

[pkgoogle]

Original Issue: tensorflow/tensorflow#60154