-
-
Notifications
You must be signed in to change notification settings - Fork 8.7k
/
test_linalg.cc
113 lines (96 loc) · 2.92 KB
/
test_linalg.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
#include <gtest/gtest.h>
#include <xgboost/host_device_vector.h>
#include <xgboost/linalg.h>
#include <numeric>
namespace xgboost {
namespace linalg {
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, size_t n_rows, size_t n_cols) {
storage->Resize(n_rows * n_cols);
auto &h_storage = storage->HostVector();
std::iota(h_storage.begin(), h_storage.end(), 0);
auto m = linalg::TensorView<float, 2>{h_storage, {n_rows, static_cast<size_t>(n_cols)}, -1};
return m;
}
TEST(Linalg, Matrix) {
size_t kRows = 31, kCols = 77;
HostDeviceVector<float> storage;
auto m = MakeMatrixFromTest(&storage, kRows, kCols);
ASSERT_EQ(m.DeviceIdx(), GenericParameter::kCpuId);
ASSERT_EQ(m(0, 0), 0);
ASSERT_EQ(m(kRows - 1, kCols - 1), storage.Size() - 1);
}
TEST(Linalg, Vector) {
size_t kRows = 31, kCols = 77;
HostDeviceVector<float> storage;
auto m = MakeMatrixFromTest(&storage, kRows, kCols);
// auto v = VectorView<float>(m, 3);
auto v = m.Slice(linalg::All(), 3);
for (size_t i = 0; i < v.Size(); ++i) {
ASSERT_EQ(v(i), m(i, 3));
}
ASSERT_EQ(v(0), 3);
}
TEST(Linalg, Tensor) {
std::vector<double> data(2 * 3 * 4, 0);
std::iota(data.begin(), data.end(), 0);
TensorView<double> t{data, {2, 3, 4}, -1};
ASSERT_EQ(t.Shape()[0], 2);
ASSERT_EQ(t.Shape()[1], 3);
ASSERT_EQ(t.Shape()[2], 4);
float v = t(0, 1, 2);
ASSERT_EQ(v, 6);
auto s = t.Slice(1, All(), All());
ASSERT_EQ(s.Shape().size(), 2);
ASSERT_EQ(s.Shape()[0], 3);
ASSERT_EQ(s.Shape()[1], 4);
std::vector<std::vector<double>> sol{
{12.0, 13.0, 14.0, 15.0}, {16.0, 17.0, 18.0, 19.0}, {20.0, 21.0, 22.0, 23.0}};
for (size_t i = 0; i < s.Shape()[0]; ++i) {
for (size_t j = 0; j < s.Shape()[1]; ++j) {
ASSERT_EQ(s(i, j), sol[i][j]);
}
}
{
// as vector
TensorView<double, 1> vec{data, {data.size()}, -1};
ASSERT_EQ(vec.Size(), data.size());
ASSERT_EQ(vec.Shape(0), data.size());
ASSERT_EQ(vec.Shape().size(), 1);
for (size_t i = 0; i < data.size(); ++i) {
ASSERT_EQ(vec(i), data[i]);
}
}
{
// as matrix
TensorView<double, 2> mat(data, {6, 4}, -1);
auto s = mat.Slice(2, All());
ASSERT_EQ(s.Shape().size(), 1);
s = mat.Slice(All(), 1);
ASSERT_EQ(s.Shape().size(), 1);
}
{
// assignment
TensorView<double, 3> t{data, {2, 3, 4}, 0};
double pi = 3.14159;
t(1, 2, 3) = pi;
ASSERT_EQ(t(1, 2, 3), pi);
}
{
// Don't assign the initial dimension, tensor should be able to deduce the correct dim
// for Slice.
TensorView<double> t{data, {2, 3, 4}, 0};
auto s = t.Slice(1, 2, All());
static_assert(decltype(s)::kDimension == 1, "");
}
}
TEST(Linalg, Empty) {
auto t = TensorView<double, 2>{{}, {0, 3}, GenericParameter::kCpuId};
for (int32_t i : {0, 1, 2}) {
auto s = t.Slice(All(), i);
ASSERT_EQ(s.Size(), 0);
ASSERT_EQ(s.Shape().size(), 1);
ASSERT_EQ(s.Shape(0), 0);
}
}
} // namespace linalg
} // namespace xgboost