def _specialized_inverse_pth_root(self, input_t, exponent, epsilon=1e-12):
input_t_f64 = tf.cast(input_t, tf.float64)
val, error = matrix_functions.inlined_matrix_inverse_pth_root(
input_t_f64,
tf.shape(input_t_f64)[0],
exponent,
iter_count=40,
ridge_epsilon=epsilon)
return tf.cast(val, tf.float32), error
def testMatrixInversePthRootFunction(self):
"""Tests for matrix inverse pth roots."""
size = 4
mat_a = np.random.rand(size, size)
mat = np.dot(mat_a, mat_a.T)
expected_mat_root = np_power(mat, -0.125)
mat_root, _ = matrix_functions.inlined_matrix_inverse_pth_root(
mat, size, -0.125, ridge_epsilon=0.0)
self.assertAllCloseAccordingToType(expected_mat_root,
mat_root,
atol=TOLERANCE,
rtol=TOLERANCE)