ddim_test.cc 2.3 KB
Newer Older
1 2 3 4
#include <sstream>
#include <vector>

#include "gtest/gtest.h"
5
#include "paddle/framework/ddim.h"
6 7 8

TEST(DDim, Equality) {
  // construct a DDim from an initialization list
9
  paddle::framework::DDim ddim = paddle::framework::make_ddim({9, 1, 5});
10 11 12 13 14
  EXPECT_EQ(ddim[0], 9);
  EXPECT_EQ(ddim[1], 1);
  EXPECT_EQ(ddim[2], 5);

  // construct a DDim from a vector
Q
qijun 已提交
15
  std::vector<int64_t> vec({9, 1, 5});
16
  paddle::framework::DDim vddim = paddle::framework::make_ddim(vec);
17 18 19 20 21 22 23
  EXPECT_EQ(ddim[0], 9);
  EXPECT_EQ(ddim[1], 1);
  EXPECT_EQ(ddim[2], 5);

  // mutate a DDim
  ddim[1] = 2;
  EXPECT_EQ(ddim[1], 2);
24 25
  paddle::framework::set(ddim, 0, 6);
  EXPECT_EQ(paddle::framework::get(ddim, 0), 6);
26 27

  // vectorize a DDim
Q
qijun 已提交
28
  std::vector<int64_t> res_vec = paddle::framework::vectorize(vddim);
29 30 31
  EXPECT_EQ(res_vec[0], 9);
  EXPECT_EQ(res_vec[1], 1);
  EXPECT_EQ(res_vec[2], 5);
32 33
  paddle::framework::Dim<3> d(3, 2, 1);
  res_vec = paddle::framework::vectorize(paddle::framework::DDim(d));
34 35 36 37 38
  EXPECT_EQ(res_vec[0], 3);
  EXPECT_EQ(res_vec[1], 2);
  EXPECT_EQ(res_vec[2], 1);

  // add two DDims
39
  paddle::framework::DDim ddim_sum = ddim + vddim;
40 41 42 43 44
  EXPECT_EQ(ddim_sum[0], 15);
  EXPECT_EQ(ddim_sum[1], 3);
  EXPECT_EQ(ddim_sum[2], 10);

  // multiply two DDims
45
  paddle::framework::DDim ddim_mul = ddim * vddim;
46 47 48 49 50
  EXPECT_EQ(ddim_mul[0], 54);
  EXPECT_EQ(ddim_mul[1], 2);
  EXPECT_EQ(ddim_mul[2], 25);

  // arity of a DDim
51
  EXPECT_EQ(paddle::framework::arity(ddim), 3);
F
fengjiayi 已提交
52
  EXPECT_EQ(ddim.size(), 3);
53 54

  // product of a DDim
55
  EXPECT_EQ(paddle::framework::product(vddim), 45);
F
fengjiayi 已提交
56 57 58
  EXPECT_EQ(
      paddle::framework::product(paddle::framework::make_ddim({3, 2, 5, 3})),
      90);
59 60 61 62

  // slice a DDim
  paddle::framework::DDim ddim2 =
      paddle::framework::make_ddim({1, 2, 3, 4, 5, 6});
F
fengjiayi 已提交
63
  paddle::framework::DDim ss = paddle::framework::slice_ddim(ddim2, 2, 5);
64 65 66 67
  EXPECT_EQ(arity(ss), 3);
  EXPECT_EQ(ss[0], 3);
  EXPECT_EQ(ss[1], 4);
  EXPECT_EQ(ss[2], 5);
F
fengjiayi 已提交
68
  paddle::framework::DDim ss2 = paddle::framework::slice_ddim(ddim2, 0, 6);
69 70 71 72 73 74 75
  EXPECT_EQ(arity(ss2), 6);
  EXPECT_EQ(ss2[0], 1);
  EXPECT_EQ(ss2[1], 2);
  EXPECT_EQ(ss2[2], 3);
  EXPECT_EQ(ss2[3], 4);
  EXPECT_EQ(ss2[4], 5);
  EXPECT_EQ(ss2[5], 6);
76 77 78 79 80
}

TEST(DDim, Print) {
  // print a DDim
  std::stringstream ss;
81
  paddle::framework::DDim ddim = paddle::framework::make_ddim({2, 3, 4});
82 83 84
  ss << ddim;
  EXPECT_EQ("2, 3, 4", ss.str());
}
反馈
建议
客服 返回
顶部