Method TestSetMatrixDiagonal

tensorflow/compiler/xla/client/lib/matrix_test.cc:91–108

template <typename T>
void MatrixTest::TestSetMatrixDiagonal() {
  XlaBuilder builder("GetMatrixDiagonal");
  Array3D<T> input(2, 3, 4);
  input.FillIota(0);
  for (const auto& kv : k_and_expected<T>()) {
    XlaOp a;
    XlaOp b;
    auto a_data = CreateR3Parameter<T>(input, 0, "a", &builder, &a);
    auto new_diag =
        CreateR2Parameter<T>(Array2D<T>{kv.second}, 1, "d", &builder, &b);

    GetMatrixDiagonal(SetMatrixDiagonal(a, b + ScalarLike(b, 1), kv.first),
                      kv.first) -
        ScalarLike(b, 1);

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

XLA_TEST_F(MatrixTest, SetMatrixDiagonal_S32) {
  TestSetMatrixDiagonal<int32>();