operator_test_util.py 文件源码

def testSqrtMatmul(self):
    # Square roots are not unique, but we should have SS^T x = Ax, and in our
    # case, we should have S = S^T, so SSx = Ax.
    with self.test_session():
      for batch_shape in [(), (
          2,
          3,)]:
        for k in [1, 4]:
          operator, mat = self._build_operator_and_mat(batch_shape, k)

          # Work with 5 simultaneous systems.  5 is arbitrary.
          x = self._rng.randn(*(batch_shape + (k, 5)))

          self._compare_results(
              expected=math_ops.matmul(mat, x).eval(),
              actual=operator.sqrt_matmul(operator.sqrt_matmul(x)))