tf.linalg.sqrtm() returns nan for invertible matrix

[itsikad]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): windows 10
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): r1.14
• Python version: 3.6.8
• CUDA/cuDNN version: No

Describe the current behavior
tf.linalg.sqrtm() returns nan, even though input matrix is invertible.
When the input matrix is constructed using numpy operations instead and then converted to tensor, the behavior is as expected.

The issue can be reproduced with the following code, both .csv's can be downloaded from: https://www.kaggle.com/gemartin/world-bank-data-1960-to-2016/downloads/world-bank-data-1960-to-2016.zip/1

import numpy as np
import pandas as pd
import tensorflow as tf


def load_word_bank_dataset():
    """
    This function loads the World Bank Data and return it as NxD numpy arrays
    """
    # Load and drop rows with missing values
    fert_rate = pd.read_csv(r'./fertility_rate.csv').dropna()
    life_exp = pd.read_csv(r'./life_expectancy.csv').dropna()

    # Keep only countries that appears in both tables
    valid_countries = list(set(fert_rate['Country Code']) & set(life_exp['Country Code']))
    fert_rate = fert_rate[fert_rate['Country Code'].isin(valid_countries)]
    life_exp = life_exp[life_exp['Country Code'].isin(valid_countries)]

    # Create numpy arrays from the years columns
    years_str_list = [str(year) for year in range(1960, 2017)]
    fert_rate = fert_rate[years_str_list].to_numpy()
    life_exp = life_exp[years_str_list].to_numpy()

    return life_exp, fert_rate

def cca_loss(x1, x2, r1=0.0, r2=0.0):
    # Extract dimensions
    N, D1 = x1.shape
    _, D2 = x2.shape
    scale_mat = np.eye(N) - 1.0 / N * np.ones((N, N))
    scale_factor = tf.cast(1. / N, dtype=tf.float32)

    # This part doesn't work
    scale_mat_tensor = tf.convert_to_tensor(scale_mat, dtype=tf.float32)
    x1_T_tensor = tf.convert_to_tensor(x1.T, dtype=tf.float32)
    h1_bar_1_tensor = tf.matmul(x1_T_tensor, scale_mat_tensor)
    cov_11_1_tensor = scale_factor * tf.matmul(h1_bar_1_tensor, h1_bar_1_tensor, adjoint_b=True)
    cov_11_1_sqrt = tf.linalg.sqrtm(cov_11_1_tensor)  # Returns nan's

    # This part works
    h1_bar_np = np.matmul(x1.T, scale_mat)
    h1_bar_2_tensor = tf.convert_to_tensor(h1_bar_np, dtype=tf.float32)
    cov_11_2_tensor = scale_factor * tf.matmul(h1_bar_2_tensor, h1_bar_2_tensor, adjoint_b=True)
    cov_11_2_sqrt = tf.linalg.sqrtm(cov_11_2_tensor)  # Returns real valued matrix

    return -1

if __name__ == '__main__':
    tf.compat.v1.enable_eager_execution()
    # Load data
    fertility_rate, life_expectancy = load_word_bank_dataset()
    loss = cca_loss(fertility_rate, life_expectancy, r1=0., r2=0.)

[oanush]

@itsikad ,
Can you please provide the complete code necessary to reproduce the issue in editable format.Thanks

[itsikad]

@oanush
Thanks for your response.
I added code to reproduce the issue.

[itsikad]

I continued investigating this issue, I believe the root cause is that the matrix is ill posed. It's minimal eigenvalue is 10^-6, the decomposition algorithm that is used to calculate its square root might get small negative values or fail to converge.
The difference between the two pieces of code is that numpy calculation are with float64 and tf with float32.
I'm closing this issue.

[kyamagu]

For anyone seeing this issue, the following approach can mitigate the nan issue though numerical stability is another issue.

tf.cast(tf.linalg.sqrtm(tf.cast(matrix, tf.complex128)), tf.float64)