如何解决Tensorflow中由低精度数据类型tf.float32导致的Cholesky分解错误？

问：

似乎 tensorflow 仅支持 tf.float32 用于训练神经网络，这会导致我的算法出现 Cholesky 分解问题。

下面的代码是我的计算图的一部分，其中 X_latent 是要传入的张量，它返回一个正定义矩阵。

def get_mat_LX(self, X_latent):
    dim_P = int(X_latent.shape[1])
    col = tf.reduce_sum(X_latent*X_latent, 1)
    col = tf.reshape(col, [-1, 1])
    prod = tf.matmul(X_latent, tf.transpose(X_latent))
    log_mat_LX = - (0.5/self.rho+0.5/self.sigma)*col + 1/self.sigma*prod - (0.5/self.rho+0.5/self.sigma)*tf.transpose(col) + 2*math.log(self.alpha**0.5)-2*dim_P*math.log((math.pi*self.rho)**0.5)
    mat_LX = tf.exp(log_mat_LX)

    return mat_LX

稍后，我需要对这个正定矩阵进行 cholesky 分解：

tensorflow_matrix = self.get_mat_LX(latent_var)
chol = tf.cholesky(tensorflow_matrix)

但是，由于数据类型问题，矩阵中的一些小值（例如 3.825485980697876e-41）将被 tensorflow 截断为 0.0000000e+00，因此 返回的矩阵不再是正定义的，并导致 cholesky 分解问题：

Cholesky decomposition was not successful. The input might not be valid. 

例如，对于

latent_var:
array([[ 1.4992752 ,  0.4027754 , -0.9438937 , ..., -1.7518392 ,
    -0.59829473, -0.9430773 ],
   [ 1.4987504 ,  1.3747126 ,  0.2517067 , ..., -0.11252454,
    -1.0589341 ,  1.9861003 ],
   [ 3.739284  ,  0.03519648, -0.7698121 , ...,  1.0399017 ,
     1.4340334 ,  0.5115541 ],
   ...,
   [ 1.8933645 , -0.8301757 , -1.9493791 , ...,  2.5583518 ,
     0.65306365,  1.7139516 ],
   [ 1.3454692 , -1.4221896 , -0.04142132, ...,  1.6484771 ,
    -2.1860263 ,  2.0473964 ],
   [ 1.3071399 , -1.4299753 ,  0.09692653, ..., -3.2884247 ,
     0.42087832,  1.6871212 ]], dtype=float32)

tensorflow_matrix = self.get_mat_LX(latent_var)

如果我们用numpy计算，

def get_numpy_mat_LX(self, latent_var):
    col = np.matrix(np.sum(latent_var*latent_var,1).reshape((-1,1)))
    prod = np.matrix(latent_var)*np.matrix(latent_var).transpose()
    log_mat_LX = - (0.5/self.rho+0.5/self.sigma)*col + 1/self.sigma*prod - (0.5/self.rho+0.5/self.sigma)*col.T + 2*math.log(self.alpha**0.5)-2*latent_var.shape[1]*math.log((math.pi*self.rho)**0.5)
    mat_LX = np.exp(log_mat_LX)

    return mat_LX

actual_matrix = self.get_numpy_mat_LX(latent_var)

我们将得到：

actual_matrix:
  [[1.99842083e-26 6.19800080e-35 3.82548598e-41 ... 1.48187120e-38
  1.40404417e-37 3.77790018e-32]
  [6.19800080e-35 4.52321870e-21 2.29967037e-41 ... 2.37359479e-34
  7.37417026e-35 6.34270132e-35]
 [3.82548598e-41 2.29967037e-41 9.54639589e-30 ... 1.57905177e-41
  1.39829079e-40 1.38537289e-44]
 ...
 [1.48187120e-38 2.37359479e-34 1.57905177e-41 ... 1.92394478e-27
  2.51794412e-36 3.13480019e-42]
 [1.40404417e-37 7.37417026e-35 1.39829079e-40 ... 2.51794412e-36
  2.30329527e-25 3.14763134e-36]
 [3.77790018e-32 6.34270132e-35 1.38537289e-44 ... 3.13480019e-42
  3.14763134e-36 2.03249691e-24]]

tensorflow_matrix:
[[1.9983998e-26 6.1980058e-35 0.0000000e+00 ... 1.4818499e-38
 1.4040427e-37 3.7778746e-32]
 [6.1980058e-35 4.5232227e-21 0.0000000e+00 ... 2.3735830e-34
  7.3741981e-35 6.3427609e-35]
 [0.0000000e+00 0.0000000e+00 9.5463307e-30 ... 0.0000000e+00
  0.0000000e+00 0.0000000e+00]
 ...
 [1.4818499e-38 2.3735830e-34 0.0000000e+00 ... 1.9239023e-27
  2.5179270e-36 0.0000000e+00]
 [1.4040427e-37 7.3741981e-35 0.0000000e+00 ... 2.5179270e-36
  2.3032796e-25 3.1476519e-36]
 [3.7778746e-32 6.3427609e-35 0.0000000e+00 ... 0.0000000e+00
  3.1476519e-36 2.0325064e-24]]

原始矩阵中的某些值被截断为 0.0000000e+00。例如

tensorflow_matrix[0,2]:
0.0
actual_matrix[0,2]:
3.825485980697876e-41

我以某种方式解决了这个问题，方法是在计算图的神经网络部分之前和之后将矩阵 tf.cast 转换为 tf.float32 和 tf.float64。 但是我想知道有没有办法解决这个问题，而不之前和之后将其更改为tf.float64。