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)