Method TestMatrixDiagonal

tensorflow/compiler/xla/client/lib/matrix_test.cc:77–88

template <typename T>
void MatrixTest::TestMatrixDiagonal() {
  XlaBuilder builder("SetMatrixDiagonal");
  Array3D<T> input(2, 3, 4);
  input.FillIota(0);
  for (const auto& kv : k_and_expected<T>()) {
    XlaOp a;
    auto a_data = CreateR3Parameter<T>(input, 0, "a", &builder, &a);
    GetMatrixDiagonal(a, kv.first);

    ComputeAndCompareR2<T>(&builder, kv.second, {a_data.get()});
  }
}

template <typename T>
void MatrixTest::TestSetMatrixDiagonal() {