PyTorch: Square root of a positive semi-definite matrix

Leockl: <p>Using PyTorch, I am wanting to work out the square root of a positive <strong>semi-definite</strong> matrix. I googled around for a PyTorch implementation but can’t seem to find the right one.</p>
<p>This is what I have found:</p>
<ul>
<li>
<a href="https://github.com/steveli/pytorch-sqrtm" rel="noopener nofollow ugc">https://github.com/steveli/pytorch-sqrtm</a> (this implementation appears to only work for positive <strong>definite</strong> matrices. I am after an implementation which works for positive <strong>semi-definite</strong> matrices).</li>
<li>
<a href="https://github.com/pytorch/pytorch/issues/25481" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/25481</a> (this implementation also appears to only work for positive <strong>definite</strong> matrices. Also, the issue is still open, so I guess they haven’t finalized a final version yet).</li>
<li>
<a href="https://github.com/msubhransu/matrix-sqrt" rel="noopener nofollow ugc">https://github.com/msubhransu/matrix-sqrt</a> (according to the 2nd Github link above, this implementation is not fully PyTorch, it uses PyTorch for backward pass and Scipy for forward pass. This implementation doesn’t say anything about positive <strong>definite</strong> or positive <strong>semi-definite</strong> matrices. Also, when I had a looked at the PyTorch code, I couldn’t understand the code because it doesn’t seem to have a class “function” to call?).</li>
</ul>
<p>Anyone know where I could find a PyTorch implementation to find the square root of a positive <strong>semi-definite</strong> matrix? Would greatly appreciate it. Many thanks in advance.</p>

KFrank: <p>Hello Leo!</p>
<aside class="quote no-group" data-username="Leockl" data-post="1" data-topic="100138" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/e9a140/48.png" class="avatar"> Leockl:</div>
<blockquote>
<p>Using PyTorch, I am wanting to work out the square root of a positive <strong>semi-definite</strong> matrix.</p>
</blockquote>
</aside>
<p>Perform the eigendecomposition of your matrix and then take the<br>
square-root of your eigenvalues.  (If any of your eigenvalues of your<br>
semi-definite matrix show up as <em>numerically</em> negative, replace them<br>
with zero.)</p>
<p>For more detail, see this post:</p>
<aside class="quote quote-modified" data-post="3" data-topic="93655">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/raising-a-tensor-to-a-fractional-power/93655/3">Raising a tensor to a fractional power</a> 
  </div>
  <blockquote>
    Hi mmutic! 

A standard approach to exponentiating a matrix is to calculate its 
eigendecomposition and then exponentiate the eigenvalues. 
Here is a pytorch version 0.3.0 script that illustrates this: 
import torch
print (torch.__version__)

m = torch.FloatTensor([[.5,.5],[.7,.9]])           # original matrix

# desired result
mres = torch.FloatTensor ([[ 2.69776664, -1.10907208], [-1.55270001,  1.81051025]])

evals, evecs = torch.eig (m, eigenvectors = True)  # get eigendecomposition
evals = e…
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

Leockl: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>, many thanks for your solution. The solution makes sense and it could definitely work for what I need to do. I am implementing a new type of classifier and one of the functions used in the square root of a positive <strong>semi-definite</strong>  matrix.</p>
<p>One quick question though, if one can do this elegantly like what you have proposed, then why are people working on a PyTorch function for this which involves more complicated math/algorithms, for eg. in this Github link below especially when you see the more recent comments from the last couple of days?</p>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/25481" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/25481" target="_blank" rel="noopener nofollow ugc">[feature request] symmetric matrix square root</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-08-30" data-time="18:19:26" data-timezone="UTC">06:19PM - 30 Aug 19 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/yaroslavvb" target="_blank" rel="noopener nofollow ugc">
          <img alt="yaroslavvb" src="https://avatars3.githubusercontent.com/u/23068?v=4" class="onebox-avatar-inline" width="20" height="20">
          yaroslavvb
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">This is needed when incorporating curvature information into optimization
There's PyTorch implementation of symmetric matrix square root op here but they use...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">function request</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: linear algebra</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


KFrank: <p>Hello Leo!</p>
<aside class="quote no-group" data-username="Leockl" data-post="3" data-topic="100138" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/l/e9a140/48.png" class="avatar"> Leockl:</div>
<blockquote>
<p>if one can do this elegantly like what you have proposed, then why are people working on a PyTorch function for this which involves more complicated math/algorithms, for eg. in this Github link below especially when you see the more recent comments from the last couple of days?</p>
</blockquote>
</aside>
<p>First comment:</p>
<p>The code Yaroslav posted at the beginning of the github issue to which<br>
you linked is basically what I suggested.  He (properly) treats the null<br>
space of the semi-definite matrix more carefully, and he (properly) uses<br>
<code>torch.symeig()</code> rather than <code>torch.eig()</code>.</p>
<p>Second comment:</p>
<p>More importantly, I’m not an expert, but I have no reason to believe that<br>
eigendecomposition is the best algorithm for the root of a matrix.  The<br>
github issue discusses other approaches that could be faster and/or<br>
numerically more stable or accurate.</p>
<p>The eigendecomposition contains, in a sense, more information than<br>
the root, so it could well be more expensive to calculate.  By way of<br>
analogy, you can use eigendecomposition to calculate the inverse of<br>
a matrix (take the reciprocals of the eigenvalues), but it is not the<br>
preferred matrix-inverse algorithm.</p>
<p>Best.</p>
<p>K. Frank</p>

zakajd: <p>For anyone interested in the same topic in the future, here is code snippets taken from <a href="https://github.com/photosynthesis-team/piq/issues/190" rel="noopener nofollow ugc">PIQ library</a></p>
<ol>
<li>Computation using iterative Newton-Schulz algorithm</li>
</ol>
<pre><code class="lang-python">def _approximation_error(matrix: torch.Tensor, s_matrix: torch.Tensor) -&gt; torch.Tensor:
    norm_of_matrix = torch.norm(matrix)
    error = matrix - torch.mm(s_matrix, s_matrix)
    error = torch.norm(error) / norm_of_matrix
    return error


def _sqrtm_newton_schulz(matrix: torch.Tensor, num_iters: int = 100) -&gt; Tuple[torch.Tensor, torch.Tensor]:
    r"""
    Square root of matrix using Newton-Schulz Iterative method
    Source: https://github.com/msubhransu/matrix-sqrt/blob/master/matrix_sqrt.py
    Args:
        matrix: matrix or batch of matrices
        num_iters: Number of iteration of the method
    Returns:
        Square root of matrix
        Error
    """
    expected_num_dims = 2
    if matrix.dim() != expected_num_dims:
        raise ValueError(f'Input dimension equals {matrix.dim()}, expected {expected_num_dims}')

    if num_iters &lt;= 0:
        raise ValueError(f'Number of iteration equals {num_iters}, expected greater than 0')

    dim = matrix.size(0)
    norm_of_matrix = matrix.norm(p='fro')
    Y = matrix.div(norm_of_matrix)
    I = torch.eye(dim, dim, requires_grad=False).to(matrix)
    Z = torch.eye(dim, dim, requires_grad=False).to(matrix)

    s_matrix = torch.empty_like(matrix)
    error = torch.empty(1).to(matrix)

    for _ in range(num_iters):
        T = 0.5 * (3.0 * I - Z.mm(Y))
        Y = Y.mm(T)
        Z = T.mm(Z)

        s_matrix = Y * torch.sqrt(norm_of_matrix)
        error = _approximation_error(matrix, s_matrix)
        if torch.isclose(error, torch.tensor([0.]).to(error), atol=1e-5):
            break

    return s_matrix, error
</code></pre>
<ol start="2">
<li>Using  eigen decomposition as discussed in this topic</li>
</ol>
<pre><code class="lang-python">def _matrix_pow(matrix: torch.Tensor, p: float) -&gt; torch.Tensor:
    r"""
    Power of a matrix using Eigen Decomposition.
    Args:
        matrix: matrix
        p: power
    Returns:
        Power of a matrix
    """
    vals, vecs = torch.eig(matrix, eigenvectors=True)
    vals = torch.view_as_complex(vals.contiguous())
    vals_pow = vals.pow(p)
    vals_pow = torch.view_as_real(vals_pow)[:, 0]
    matrix_pow = torch.matmul(vecs, torch.matmul(torch.diag(vals_pow), torch.inverse(vecs)))
    return matrix_pow
</code></pre>