1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/util/sparse/sparse_tensor.h"
17
18 #include <string>
19 #include <vector>
20
21 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
22 #include "tensorflow/core/framework/tensor.h"
23 #include "tensorflow/core/framework/tensor_types.h"
24 #include "tensorflow/core/lib/core/status_test_util.h"
25 #include "tensorflow/core/lib/random/simple_philox.h"
26 #include "tensorflow/core/lib/strings/str_util.h"
27 #include "tensorflow/core/platform/test.h"
28 #include "tensorflow/core/platform/test_benchmark.h"
29
30 namespace tensorflow {
31 namespace sparse {
32 namespace {
33
34 Eigen::Tensor<int64, 2, Eigen::RowMajor, Eigen::DenseIndex>
GetSimpleIndexTensor(int N,const int NDIM)35 GetSimpleIndexTensor(int N, const int NDIM) {
36 Eigen::Tensor<int64, 2, Eigen::RowMajor, Eigen::DenseIndex> ix(N, NDIM);
37 ix(0, 0) = 0;
38 ix(0, 1) = 0;
39 ix(0, 2) = 0;
40
41 ix(1, 0) = 3;
42 ix(1, 1) = 0;
43 ix(1, 2) = 0;
44
45 ix(2, 0) = 2;
46 ix(2, 1) = 0;
47 ix(2, 2) = 0;
48
49 ix(3, 0) = 0;
50 ix(3, 1) = 1;
51 ix(3, 2) = 0;
52
53 ix(4, 0) = 0;
54 ix(4, 1) = 0;
55 ix(4, 2) = 2;
56 return ix;
57 }
58
TEST(SparseTensorTest,DimComparatorSorts)59 TEST(SparseTensorTest, DimComparatorSorts) {
60 int64_t N = 5;
61 const int NDIM = 3;
62 auto ix = GetSimpleIndexTensor(N, NDIM);
63 TTypes<int64>::Matrix map(ix.data(), N, NDIM);
64
65 std::vector<int64> sorting(N);
66 for (std::size_t n = 0; n < N; ++n) sorting[n] = n;
67
68 // new order should be: {0, 4, 3, 2, 1}
69 std::vector<int64> order{0, 1, 2};
70 std::vector<int64> shape{N, N, N};
71 DimComparator sorter(map, order, shape);
72 std::sort(sorting.begin(), sorting.end(), sorter);
73 EXPECT_EQ(sorting, std::vector<int64>({0, 4, 3, 2, 1}));
74
75 FixedDimComparator<3> sorter_fixed(map, order, shape);
76 std::sort(sorting.begin(), sorting.end(), sorter_fixed);
77 EXPECT_EQ(sorting, std::vector<int64>({0, 4, 3, 2, 1}));
78
79 // new order should be: {0, 3, 2, 1, 4}
80 std::vector<int64> order1{2, 0, 1};
81 DimComparator sorter1(map, order1, shape);
82 for (std::size_t n = 0; n < N; ++n) sorting[n] = n;
83 std::sort(sorting.begin(), sorting.end(), sorter1);
84 EXPECT_EQ(sorting, std::vector<int64>({0, 3, 2, 1, 4}));
85
86 FixedDimComparator<3> sorter1_fixed(map, order1, shape);
87 for (std::size_t n = 0; n < N; ++n) sorting[n] = n;
88 std::sort(sorting.begin(), sorting.end(), sorter1_fixed);
89 EXPECT_EQ(sorting, std::vector<int64>({0, 3, 2, 1, 4}));
90 }
91
TEST(SparseTensorTest,SparseTensorInvalidIndicesType)92 TEST(SparseTensorTest, SparseTensorInvalidIndicesType) {
93 int N = 5;
94 const int NDIM = 3;
95 Tensor ix(DT_INT32, TensorShape({N, NDIM}));
96 Tensor vals(DT_STRING, TensorShape({N}));
97 SparseTensor result;
98
99 EXPECT_EQ(SparseTensor::Create(ix, vals, TensorShape({10, 10, 10}), {0, 1, 2},
100 &result)
101 .code(),
102 error::INVALID_ARGUMENT);
103 }
104
TEST(SparseTensorTest,SparseTensorInvalidIndicesShape)105 TEST(SparseTensorTest, SparseTensorInvalidIndicesShape) {
106 int N = 5;
107 const int NDIM = 3;
108 Tensor ix(DT_INT64, TensorShape({N, NDIM, 1}));
109 Tensor vals(DT_STRING, TensorShape({N}));
110 SparseTensor result;
111
112 EXPECT_EQ(SparseTensor::Create(ix, vals, TensorShape({10, 10, 10}), {0, 1, 2},
113 &result)
114 .code(),
115 error::INVALID_ARGUMENT);
116 }
117
TEST(SparseTensorTest,SparseTensorInvalidValues)118 TEST(SparseTensorTest, SparseTensorInvalidValues) {
119 int N = 5;
120 const int NDIM = 3;
121 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
122 Tensor vals(DT_STRING, TensorShape({N, 1}));
123 SparseTensor result;
124
125 EXPECT_EQ(SparseTensor::Create(ix, vals, TensorShape({10, 10, 10}), {0, 1, 2},
126 &result)
127 .code(),
128 error::INVALID_ARGUMENT);
129 }
130
TEST(SparseTensorTest,SparseTensorInvalidN)131 TEST(SparseTensorTest, SparseTensorInvalidN) {
132 int N = 5;
133 const int NDIM = 3;
134 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
135 Tensor vals(DT_STRING, TensorShape({N - 1}));
136 SparseTensor result;
137
138 EXPECT_EQ(SparseTensor::Create(ix, vals, TensorShape({10, 10, 10}), {0, 1, 2},
139 &result)
140 .code(),
141 error::INVALID_ARGUMENT);
142 }
143
TEST(SparseTensorTest,SparseTensorInvalidOrder)144 TEST(SparseTensorTest, SparseTensorInvalidOrder) {
145 int N = 5;
146 const int NDIM = 3;
147 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
148 Tensor vals(DT_STRING, TensorShape({N}));
149 SparseTensor result;
150
151 EXPECT_EQ(
152 SparseTensor::Create(ix, vals, TensorShape({10, 10, 10}), {0, 1}, &result)
153 .code(),
154 error::INVALID_ARGUMENT);
155 }
TEST(SparseTensorTest,SparseTensorInvalidShape)156 TEST(SparseTensorTest, SparseTensorInvalidShape) {
157 int N = 5;
158 const int NDIM = 3;
159 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
160 Tensor vals(DT_STRING, TensorShape({N}));
161 SparseTensor result;
162
163 EXPECT_EQ(
164 SparseTensor::Create(ix, vals, TensorShape({10, 10}), {0, 1, 2}, &result)
165 .code(),
166 error::INVALID_ARGUMENT);
167 }
168
TEST(SparseTensorTest,SparseTensorConstruction)169 TEST(SparseTensorTest, SparseTensorConstruction) {
170 int N = 5;
171 const int NDIM = 3;
172 auto ix_c = GetSimpleIndexTensor(N, NDIM);
173 Eigen::Tensor<tstring, 1, Eigen::RowMajor> vals_c(N);
174 vals_c(0) = "hi0";
175 vals_c(1) = "hi1";
176 vals_c(2) = "hi2";
177 vals_c(3) = "hi3";
178 vals_c(4) = "hi4";
179
180 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
181 Tensor vals(DT_STRING, TensorShape({N}));
182
183 auto ix_t = ix.matrix<int64>();
184 auto vals_t = vals.vec<tstring>();
185 vals_t = vals_c;
186 ix_t = ix_c;
187
188 TensorShape shape({10, 10, 10});
189 std::vector<int64> order{0, 1, 2};
190 SparseTensor st;
191 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
192 Status st_indices_valid = st.IndicesValid();
193 EXPECT_FALSE(st_indices_valid.ok());
194 EXPECT_EQ(
195 "indices[2] = [2,0,0] is out of order. "
196 "Many sparse ops require sorted indices.\n"
197 " Use `tf.sparse.reorder` to create a correctly ordered copy."
198 "\n\n",
199 st_indices_valid.error_message());
200
201 // Regardless of how order is updated; so long as there are no
202 // duplicates, the resulting indices are valid.
203 st.Reorder<tstring>({2, 0, 1});
204 TF_EXPECT_OK(st.IndicesValid());
205 EXPECT_EQ(vals_t(0), "hi0");
206 EXPECT_EQ(vals_t(1), "hi3");
207 EXPECT_EQ(vals_t(2), "hi2");
208 EXPECT_EQ(vals_t(3), "hi1");
209 EXPECT_EQ(vals_t(4), "hi4");
210
211 ix_t = ix_c;
212 vals_t = vals_c;
213 st.Reorder<tstring>({0, 1, 2});
214 TF_EXPECT_OK(st.IndicesValid());
215 EXPECT_EQ(vals_t(0), "hi0");
216 EXPECT_EQ(vals_t(1), "hi4");
217 EXPECT_EQ(vals_t(2), "hi3");
218 EXPECT_EQ(vals_t(3), "hi2");
219 EXPECT_EQ(vals_t(4), "hi1");
220
221 ix_t = ix_c;
222 vals_t = vals_c;
223 st.Reorder<tstring>({2, 1, 0});
224 TF_EXPECT_OK(st.IndicesValid());
225 }
226
TEST(SparseTensorTest,EmptySparseTensorAllowed)227 TEST(SparseTensorTest, EmptySparseTensorAllowed) {
228 int N = 0;
229 const int NDIM = 3;
230
231 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
232 Tensor vals(DT_STRING, TensorShape({N}));
233
234 std::vector<int64> shape{10, 10, 10};
235 std::vector<int64> order{0, 1, 2};
236 SparseTensor st;
237 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
238 TF_EXPECT_OK(st.IndicesValid());
239 EXPECT_EQ(st.order(), order);
240
241 std::vector<int64> new_order{1, 0, 2};
242 st.Reorder<tstring>(new_order);
243 TF_EXPECT_OK(st.IndicesValid());
244 EXPECT_EQ(st.order(), new_order);
245 }
246
TEST(SparseTensorTest,SortingWorksCorrectly)247 TEST(SparseTensorTest, SortingWorksCorrectly) {
248 int N = 30;
249 const int NDIM = 4;
250
251 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
252 Tensor vals(DT_STRING, TensorShape({N}));
253 TensorShape shape({1000, 1000, 1000, 1000});
254 SparseTensor st;
255 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, &st));
256
257 auto ix_t = ix.matrix<int64>();
258
259 for (int n = 0; n < 100; ++n) {
260 ix_t = ix_t.random(Eigen::internal::UniformRandomGenerator<int64>(n + 1));
261 ix_t = ix_t.abs() % 1000;
262 st.Reorder<tstring>({0, 1, 2, 3});
263 TF_EXPECT_OK(st.IndicesValid());
264 st.Reorder<tstring>({3, 2, 1, 0});
265 TF_EXPECT_OK(st.IndicesValid());
266 st.Reorder<tstring>({1, 0, 2, 3});
267 TF_EXPECT_OK(st.IndicesValid());
268 st.Reorder<tstring>({3, 0, 2, 1});
269 TF_EXPECT_OK(st.IndicesValid());
270 }
271 }
272
TEST(SparseTensorTest,ValidateIndicesFindsInvalid)273 TEST(SparseTensorTest, ValidateIndicesFindsInvalid) {
274 int N = 2;
275 const int NDIM = 3;
276
277 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
278 Tensor vals(DT_STRING, TensorShape({N}));
279
280 Eigen::Tensor<int64, 2, Eigen::RowMajor> ix_orig(N, NDIM);
281 ix_orig(0, 0) = 0;
282 ix_orig(0, 1) = 0;
283 ix_orig(0, 2) = 0;
284
285 ix_orig(1, 0) = 0;
286 ix_orig(1, 1) = 0;
287 ix_orig(1, 2) = 0;
288
289 auto ix_t = ix.matrix<int64>();
290 ix_t = ix_orig;
291
292 TensorShape shape({10, 10, 10});
293 std::vector<int64> order{0, 1, 2};
294 SparseTensor st;
295 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
296
297 st.Reorder<tstring>(order);
298 Status st_indices_valid = st.IndicesValid();
299 EXPECT_FALSE(st_indices_valid.ok());
300 EXPECT_EQ("indices[1] = [0,0,0] is repeated",
301 st_indices_valid.error_message());
302
303 ix_orig(1, 2) = 1;
304 ix_t = ix_orig;
305 st.Reorder<tstring>(order);
306 TF_EXPECT_OK(st.IndicesValid()); // second index now (0, 0, 1)
307
308 ix_orig(0, 2) = 1;
309 ix_t = ix_orig;
310 st.Reorder<tstring>(order);
311 st_indices_valid = st.IndicesValid();
312 EXPECT_FALSE(st_indices_valid.ok()); // first index now (0, 0, 1)
313 EXPECT_EQ("indices[1] = [0,0,1] is repeated",
314 st_indices_valid.error_message());
315 }
316
TEST(SparseTensorTest,SparseTensorCheckBoundaries)317 TEST(SparseTensorTest, SparseTensorCheckBoundaries) {
318 int N = 5;
319 const int NDIM = 3;
320
321 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
322 Tensor vals(DT_STRING, TensorShape({N}));
323
324 auto ix_t = GetSimpleIndexTensor(N, NDIM);
325
326 ix.matrix<int64>() = ix_t;
327
328 TensorShape shape({10, 10, 10});
329 std::vector<int64> order{0, 1, 2};
330
331 SparseTensor st;
332 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
333 EXPECT_FALSE(st.IndicesValid().ok());
334
335 st.Reorder<tstring>(order);
336 TF_EXPECT_OK(st.IndicesValid());
337
338 ix_t(0, 0) = 11;
339 ix.matrix<int64>() = ix_t;
340 st.Reorder<tstring>(order);
341 Status st_indices_valid = st.IndicesValid();
342 EXPECT_FALSE(st_indices_valid.ok());
343 // Error message references index 4 because of the call to Reorder.
344 EXPECT_EQ("[11,0,0] is out of bounds: need 0 <= index < [10,10,10]",
345 st_indices_valid.error_message().substr(13));
346
347 ix_t(0, 0) = -1;
348 ix.matrix<int64>() = ix_t;
349 st.Reorder<tstring>(order);
350 st_indices_valid = st.IndicesValid();
351 EXPECT_FALSE(st_indices_valid.ok());
352 EXPECT_EQ("[-1,0,0] is out of bounds: need 0 <= index < [10,10,10]",
353 st_indices_valid.error_message().substr(13));
354
355 ix_t(0, 0) = 0;
356 ix.matrix<int64>() = ix_t;
357 st.Reorder<tstring>(order);
358 TF_EXPECT_OK(st.IndicesValid());
359 }
360
TEST(SparseTensorTest,SparseTensorToDenseTensor)361 TEST(SparseTensorTest, SparseTensorToDenseTensor) {
362 int N = 5;
363 const int NDIM = 3;
364
365 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
366 Tensor vals(DT_STRING, TensorShape({N}));
367
368 auto ix_t = GetSimpleIndexTensor(N, NDIM);
369 auto vals_t = vals.vec<tstring>();
370
371 ix.matrix<int64>() = ix_t;
372
373 vals_t(0) = "hi0";
374 vals_t(1) = "hi1";
375 vals_t(2) = "hi2";
376 vals_t(3) = "hi3";
377 vals_t(4) = "hi4";
378
379 TensorShape shape({4, 4, 5});
380 std::vector<int64> order{0, 1, 2};
381 SparseTensor st;
382 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
383
384 Tensor dense(DT_STRING, TensorShape({4, 4, 5}));
385 st.ToDense<tstring>(&dense);
386
387 auto dense_t = dense.tensor<tstring, 3>();
388 Eigen::array<Eigen::DenseIndex, NDIM> ix_n;
389 for (int n = 0; n < N; ++n) {
390 for (int d = 0; d < NDIM; ++d) ix_n[d] = ix_t(n, d);
391 EXPECT_EQ(dense_t(ix_n), vals_t(n));
392 }
393
394 // Spot checks on the others
395 EXPECT_EQ(dense_t(0, 0, 1), "");
396 EXPECT_EQ(dense_t(0, 0, 3), "");
397 EXPECT_EQ(dense_t(3, 3, 3), "");
398 EXPECT_EQ(dense_t(3, 3, 4), "");
399 }
400
TEST(SparseTensorTest,SparseTensorToLargerDenseTensor)401 TEST(SparseTensorTest, SparseTensorToLargerDenseTensor) {
402 int N = 5;
403 const int NDIM = 3;
404
405 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
406 Tensor vals(DT_STRING, TensorShape({N}));
407
408 auto ix_t = GetSimpleIndexTensor(N, NDIM);
409 auto vals_t = vals.vec<tstring>();
410
411 ix.matrix<int64>() = ix_t;
412
413 vals_t(0) = "hi0";
414 vals_t(1) = "hi1";
415 vals_t(2) = "hi2";
416 vals_t(3) = "hi3";
417 vals_t(4) = "hi4";
418
419 TensorShape shape({4, 4, 5});
420 std::vector<int64> order{0, 1, 2};
421 SparseTensor st;
422 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
423
424 Tensor dense(DT_STRING, TensorShape({10, 10, 10}));
425 st.ToDense<tstring>(&dense);
426
427 auto dense_t = dense.tensor<tstring, 3>();
428 Eigen::array<Eigen::DenseIndex, NDIM> ix_n;
429 for (int n = 0; n < N; ++n) {
430 for (int d = 0; d < NDIM; ++d) ix_n[d] = ix_t(n, d);
431 EXPECT_EQ(dense_t(ix_n), vals_t(n));
432 }
433
434 // Spot checks on the others
435 EXPECT_EQ(dense_t(0, 0, 1), "");
436 EXPECT_EQ(dense_t(0, 0, 3), "");
437 EXPECT_EQ(dense_t(3, 3, 3), "");
438 EXPECT_EQ(dense_t(3, 3, 4), "");
439 EXPECT_EQ(dense_t(9, 0, 0), "");
440 EXPECT_EQ(dense_t(9, 0, 9), "");
441 EXPECT_EQ(dense_t(9, 9, 9), "");
442 }
443
TEST(SparseTensorTest,SparseTensorGroup)444 TEST(SparseTensorTest, SparseTensorGroup) {
445 int N = 5;
446 const int NDIM = 3;
447
448 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
449 Tensor vals(DT_INT32, TensorShape({N}));
450
451 auto ix_t = ix.matrix<int64>();
452 auto vals_t = vals.vec<int32>();
453
454 ix_t = GetSimpleIndexTensor(N, NDIM);
455
456 vals_t(0) = 1; // associated with ix (000)
457 vals_t(1) = 2; // associated with ix (300)
458 vals_t(2) = 3; // associated with ix (200)
459 vals_t(3) = 4; // associated with ix (010)
460 vals_t(4) = 5; // associated with ix (002)
461
462 TensorShape shape({10, 10, 10});
463 std::vector<int64> order{0, 1, 2};
464
465 SparseTensor st;
466 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
467 st.Reorder<int32>(order);
468
469 std::vector<std::vector<int64> > groups;
470 std::vector<TTypes<int64>::UnalignedConstMatrix> grouped_indices;
471 std::vector<TTypes<int32>::UnalignedVec> grouped_values;
472
473 // Group by index 0
474 auto gi = st.group({0});
475
476 // All the hard work is right here!
477 for (const auto& g : gi) {
478 groups.push_back(g.group());
479 VLOG(1) << "Group: " << absl::StrJoin(g.group(), ",");
480 VLOG(1) << "Indices: " << g.indices();
481 VLOG(1) << "Values: " << g.values<int32>();
482
483 grouped_indices.push_back(g.indices());
484 grouped_values.push_back(g.values<int32>());
485 }
486
487 // Group by dimension 0, we have groups: 0--, 2--, 3--
488 EXPECT_EQ(groups.size(), 3);
489 EXPECT_EQ(groups[0], std::vector<int64>({0}));
490 EXPECT_EQ(groups[1], std::vector<int64>({2}));
491 EXPECT_EQ(groups[2], std::vector<int64>({3}));
492
493 std::vector<Eigen::Tensor<int64, 2, Eigen::RowMajor> > expected_indices;
494 std::vector<Eigen::Tensor<int32, 1, Eigen::RowMajor> > expected_vals;
495
496 // First group: 000, 002, 010
497 expected_indices.emplace_back(3, NDIM); // 3 x 3 tensor
498 expected_vals.emplace_back(3); // 3 x 5 x 1 x 1 tensor
499 expected_indices[0].setZero();
500 expected_indices[0](1, 2) = 2; // 002
501 expected_indices[0](2, 1) = 1; // 010
502 expected_vals[0].setConstant(-1);
503 expected_vals[0](0) = 1; // val associated with ix 000
504 expected_vals[0](1) = 5; // val associated with ix 002
505 expected_vals[0](2) = 4; // val associated with ix 010
506
507 // Second group: 200
508 expected_indices.emplace_back(1, NDIM);
509 expected_vals.emplace_back(1);
510 expected_indices[1].setZero();
511 expected_indices[1](0, 0) = 2; // 200
512 expected_vals[1](0) = 3; // val associated with ix 200
513
514 // Third group: 300
515 expected_indices.emplace_back(1, NDIM);
516 expected_vals.emplace_back(1);
517 expected_indices[2].setZero();
518 expected_indices[2](0, 0) = 3; // 300
519 expected_vals[2](0) = 2; // val associated with ix 300
520
521 for (std::size_t gix = 0; gix < groups.size(); ++gix) {
522 // Compare indices
523 auto gi_t = grouped_indices[gix];
524 Eigen::Tensor<bool, 0, Eigen::RowMajor> eval =
525 (gi_t == expected_indices[gix]).all();
526 EXPECT_TRUE(eval()) << gix << " indices: " << gi_t << " vs. "
527 << expected_indices[gix];
528
529 // Compare values
530 auto gv_t = grouped_values[gix];
531 eval = (gv_t == expected_vals[gix]).all();
532 EXPECT_TRUE(eval()) << gix << " values: " << gv_t << " vs. "
533 << expected_vals[gix];
534 }
535 }
536
TEST(SparseTensorTest,Concat)537 TEST(SparseTensorTest, Concat) {
538 int N = 5;
539 const int NDIM = 3;
540
541 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
542 Tensor vals(DT_STRING, TensorShape({N}));
543
544 auto ix_c = GetSimpleIndexTensor(N, NDIM);
545
546 auto ix_t = ix.matrix<int64>();
547 auto vals_t = vals.vec<tstring>();
548
549 ix_t = ix_c;
550
551 TensorShape shape({10, 10, 10});
552 std::vector<int64> order{0, 1, 2};
553
554 SparseTensor st;
555 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
556 EXPECT_FALSE(st.IndicesValid().ok());
557 st.Reorder<tstring>(order);
558 TF_EXPECT_OK(st.IndicesValid());
559
560 SparseTensor concatted = SparseTensor::Concat<tstring>({st, st, st, st});
561 EXPECT_EQ(concatted.order(), st.order());
562 gtl::InlinedVector<int64, 8> expected_shape{40, 10, 10};
563 EXPECT_EQ(concatted.shape(), expected_shape);
564 EXPECT_EQ(concatted.num_entries(), 4 * N);
565 TF_EXPECT_OK(concatted.IndicesValid());
566
567 auto conc_ix_t = concatted.indices().matrix<int64>();
568 auto conc_vals_t = concatted.values().vec<tstring>();
569
570 for (int n = 0; n < 4; ++n) {
571 for (int i = 0; i < N; ++i) {
572 // Dimensions match except the primary dim, which is offset by
573 // shape[order[0]]
574 EXPECT_EQ(conc_ix_t(n * N + i, 0), 10 * n + ix_t(i, 0));
575 EXPECT_EQ(conc_ix_t(n * N + i, 1), ix_t(i, 1));
576 EXPECT_EQ(conc_ix_t(n * N + i, 1), ix_t(i, 1));
577
578 // Values match
579 EXPECT_EQ(conc_vals_t(n * N + i), vals_t(i));
580 }
581 }
582
583 // Concat works if non-primary ix is out of order, but output order
584 // is not defined
585 SparseTensor st_ooo;
586 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, {0, 2, 1},
587 &st_ooo)); // non-primary ix OOO
588 SparseTensor conc_ooo = SparseTensor::Concat<tstring>({st, st, st, st_ooo});
589 std::vector<int64> expected_ooo{-1, -1, -1};
590 EXPECT_EQ(conc_ooo.order(), expected_ooo);
591 EXPECT_EQ(conc_ooo.shape(), expected_shape);
592 EXPECT_EQ(conc_ooo.num_entries(), 4 * N);
593 }
594
TEST(SparseTensorTest,ConcatEmptyN)595 TEST(SparseTensorTest, ConcatEmptyN) {
596 constexpr int N = 0;
597 constexpr int NDIM = 2;
598 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
599 Tensor vals(DT_STRING, TensorShape({N}));
600 TensorShape shape({10, 10});
601 SparseTensor st;
602 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, {0, 1}, &st));
603
604 SparseTensor concatted = SparseTensor::Concat<tstring>({st, st, st});
605
606 EXPECT_EQ(concatted.num_entries(), 0);
607 }
608
609 // TODO(ebrevdo): ReduceToDense(R={dim1,dim2,...}, reduce_fn, &output)
610 // reduce_fn sees slices of resorted values based on generator (dim: DDIMS), and
611 // slices of resorted indices on generator.
612
TEST(SparseTensorTest,Split)613 TEST(SparseTensorTest, Split) {
614 const int N = 4;
615 const int DIM = 2;
616
617 Tensor ids(DT_INT64, TensorShape({N, DIM}));
618 Tensor vals(DT_INT64, TensorShape({N}));
619
620 ids.matrix<int64>()(0, 0) = 0;
621 ids.matrix<int64>()(0, 1) = 0;
622 ids.matrix<int64>()(1, 0) = 1;
623 ids.matrix<int64>()(1, 1) = 1;
624 ids.matrix<int64>()(2, 0) = 1;
625 ids.matrix<int64>()(2, 1) = 2;
626 ids.matrix<int64>()(3, 0) = 3;
627 ids.matrix<int64>()(3, 1) = 0;
628
629 vals.vec<int64>()(0) = 1;
630 vals.vec<int64>()(1) = 2;
631 vals.vec<int64>()(2) = 3;
632 vals.vec<int64>()(3) = 4;
633
634 SparseTensor st;
635 TF_ASSERT_OK(SparseTensor::Create(ids, vals, TensorShape({4, 3}), &st));
636
637 std::vector<SparseTensor> st_list;
638 TF_ASSERT_OK(SparseTensor::Split<int64>(st, 0, 2, &st_list));
639
640 EXPECT_EQ(st_list.size(), 2);
641 auto expected_shape = gtl::InlinedVector<int64, 8>{2, 3};
642
643 EXPECT_EQ(st_list[0].shape(), expected_shape);
644 EXPECT_EQ(st_list[0].values().NumElements(), 3);
645 EXPECT_EQ(st_list[0].values().vec<int64>()(0), 1);
646 EXPECT_EQ(st_list[0].values().vec<int64>()(1), 2);
647 EXPECT_EQ(st_list[0].values().vec<int64>()(2), 3);
648 EXPECT_EQ(st_list[0].indices().NumElements(), 6);
649 EXPECT_EQ(st_list[0].indices().matrix<int64>()(0, 0), 0);
650 EXPECT_EQ(st_list[0].indices().matrix<int64>()(0, 1), 0);
651 EXPECT_EQ(st_list[0].indices().matrix<int64>()(1, 0), 1);
652 EXPECT_EQ(st_list[0].indices().matrix<int64>()(1, 1), 1);
653 EXPECT_EQ(st_list[0].indices().matrix<int64>()(2, 0), 1);
654 EXPECT_EQ(st_list[0].indices().matrix<int64>()(2, 1), 2);
655
656 EXPECT_EQ(st_list[1].shape(), expected_shape);
657 EXPECT_EQ(st_list[1].values().NumElements(), 1);
658 EXPECT_EQ(st_list[1].values().vec<int64>()(0), 4);
659 EXPECT_EQ(st_list[1].indices().NumElements(), 2);
660 EXPECT_EQ(st_list[1].indices().matrix<int64>()(0, 0), 1);
661 EXPECT_EQ(st_list[1].indices().matrix<int64>()(0, 1), 0);
662 }
663
TEST(SparseTensorTest,Slice)664 TEST(SparseTensorTest, Slice) {
665 const int N = 4;
666 const int DIM = 2;
667
668 Tensor ids(DT_INT64, TensorShape({N, DIM}));
669 Tensor vals(DT_INT64, TensorShape({N}));
670
671 ids.matrix<int64>()(0, 0) = 0;
672 ids.matrix<int64>()(0, 1) = 0;
673 ids.matrix<int64>()(1, 0) = 1;
674 ids.matrix<int64>()(1, 1) = 1;
675 ids.matrix<int64>()(2, 0) = 1;
676 ids.matrix<int64>()(2, 1) = 2;
677 ids.matrix<int64>()(3, 0) = 3;
678 ids.matrix<int64>()(3, 1) = 0;
679
680 vals.vec<int64>()(0) = 1;
681 vals.vec<int64>()(1) = 2;
682 vals.vec<int64>()(2) = 3;
683 vals.vec<int64>()(3) = 4;
684
685 SparseTensor st;
686 TF_ASSERT_OK(SparseTensor::Create(ids, vals, TensorShape({4, 3}), &st));
687
688 std::vector<int64> start(2, 0);
689 std::vector<int64> size(2);
690 size[0] = 2;
691 size[1] = 3;
692
693 SparseTensor slice = SparseTensor::Slice<int64>(st, start, size);
694
695 EXPECT_EQ(TensorShape(slice.shape()), TensorShape({2, 3}));
696 EXPECT_EQ(slice.values().NumElements(), 3);
697 EXPECT_EQ(slice.values().vec<int64>()(0), 1);
698 EXPECT_EQ(slice.values().vec<int64>()(1), 2);
699 EXPECT_EQ(slice.values().vec<int64>()(2), 3);
700 EXPECT_EQ(slice.indices().NumElements(), 6);
701 EXPECT_EQ(slice.indices().matrix<int64>()(0, 0), 0);
702 EXPECT_EQ(slice.indices().matrix<int64>()(0, 1), 0);
703 EXPECT_EQ(slice.indices().matrix<int64>()(1, 0), 1);
704 EXPECT_EQ(slice.indices().matrix<int64>()(1, 1), 1);
705 EXPECT_EQ(slice.indices().matrix<int64>()(2, 0), 1);
706 EXPECT_EQ(slice.indices().matrix<int64>()(2, 1), 2);
707 }
708
TEST(SparseTensorTest,SliceReducesOutputDimension)709 TEST(SparseTensorTest, SliceReducesOutputDimension) {
710 const int num_rows = 2;
711 const int num_columns = 2;
712
713 Tensor ids(DT_INT64, TensorShape({num_rows, num_columns}));
714 ids.matrix<int64>()(0, 0) = 0;
715 ids.matrix<int64>()(0, 1) = 0;
716 ids.matrix<int64>()(1, 0) = 1;
717 ids.matrix<int64>()(1, 1) = 1;
718
719 Tensor vals(DT_INT64, TensorShape({2}));
720 vals.vec<int64>()(0) = 1;
721 vals.vec<int64>()(1) = 2;
722
723 SparseTensor st;
724 TF_ASSERT_OK(SparseTensor::Create(ids, vals,
725 TensorShape({num_rows, num_columns}), &st));
726
727 SparseTensor slice =
728 SparseTensor::Slice<int64>(st, {num_rows + 1, 1}, {1, num_columns});
729 EXPECT_EQ(TensorShape(slice.shape()), TensorShape({0, 1}));
730 }
731
TEST(SparseTensorTest,Dim0SparseTensorToDenseTensor)732 TEST(SparseTensorTest, Dim0SparseTensorToDenseTensor) {
733 Tensor ix(DT_INT64, TensorShape({1, 0}));
734 Tensor vals(DT_INT32, TensorShape({1}));
735 vals.scalar<int32>()() = 5;
736
737 TensorShape shape({});
738 SparseTensor st;
739 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, &st));
740
741 Tensor dense(DT_INT32, TensorShape({}));
742 st.ToDense<int32>(&dense);
743
744 EXPECT_EQ(dense.scalar<int32>()(), 5);
745 }
746
BM_SparseReorderFloat(::testing::benchmark::State & state)747 static void BM_SparseReorderFloat(::testing::benchmark::State& state) {
748 int N32 = state.range(0);
749 int NDIM32 = state.range(1);
750 random::PhiloxRandom philox(301, 17);
751 random::SimplePhilox rnd(&philox);
752 const int64_t NDIM = static_cast<int64>(NDIM32);
753 const int64_t N = static_cast<int64>(N32);
754 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
755 Tensor vals(DT_FLOAT, TensorShape({N}));
756 TensorShape shape;
757 std::vector<int64> order;
758 for (int d = 0; d < NDIM32; ++d) {
759 shape.AddDim(1000);
760 order.push_back(d);
761 }
762 std::vector<int64> reorder;
763 reorder.push_back(1);
764 reorder.push_back(0);
765 for (int d = 2; d < NDIM32; ++d) {
766 reorder.push_back(d);
767 }
768 auto ix_t = ix.matrix<int64>();
769
770 for (auto s : state) {
771 state.PauseTiming();
772 for (int64_t i = 0; i < N; ++i) {
773 for (int d = 0; d < NDIM32; ++d) {
774 ix_t(i, d) = rnd.Rand64() % 1000;
775 }
776 }
777 SparseTensor st;
778 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
779
780 state.ResumeTiming();
781 st.Reorder<float>(reorder);
782 }
783 }
784
BM_SparseReorderString(::testing::benchmark::State & state)785 static void BM_SparseReorderString(::testing::benchmark::State& state) {
786 int N32 = state.range(0);
787 int NDIM32 = state.range(1);
788 random::PhiloxRandom philox(301, 17);
789 random::SimplePhilox rnd(&philox);
790 const int64_t NDIM = static_cast<int64>(NDIM32);
791 const int64_t N = static_cast<int64>(N32);
792 Tensor ix(DT_INT64, TensorShape({N, NDIM}));
793 Tensor vals(DT_STRING, TensorShape({N}));
794 TensorShape shape;
795 std::vector<int64> order;
796 auto ix_t = ix.matrix<int64>();
797 auto vals_t = vals.vec<tstring>();
798 for (int i = 0; i < N32; ++i) {
799 int len = rnd.Rand32() % 1000;
800 vals_t(i).resize(len);
801 }
802 for (int d = 0; d < NDIM32; ++d) {
803 shape.AddDim(1000);
804 order.push_back(d);
805 }
806 std::vector<int64> reorder;
807 reorder.push_back(1);
808 reorder.push_back(0);
809 for (int d = 2; d < NDIM32; ++d) {
810 reorder.push_back(d);
811 }
812
813 for (auto s : state) {
814 state.PauseTiming();
815 for (int64_t i = 0; i < N; ++i) {
816 for (int d = 0; d < NDIM32; ++d) {
817 ix_t(i, d) = rnd.Rand64() % 1000;
818 }
819 }
820 SparseTensor st;
821 TF_ASSERT_OK(SparseTensor::Create(ix, vals, shape, order, &st));
822
823 state.ResumeTiming();
824 st.Reorder<tstring>(reorder);
825 }
826 }
827
828 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(10, 2);
829 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(100, 2);
830 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(1000, 2);
831 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(10000, 2);
832 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(100000, 2);
833 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(10, 3);
834 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(100, 3);
835 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(1000, 3);
836 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(10000, 3);
837 BENCHMARK(BM_SparseReorderFloat)->UseRealTime()->ArgPair(100000, 3);
838
839 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(10, 2);
840 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(100, 2);
841 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(1000, 2);
842 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(10000, 2);
843 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(10, 3);
844 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(100, 3);
845 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(1000, 3);
846 BENCHMARK(BM_SparseReorderString)->UseRealTime()->ArgPair(10000, 3);
847
848 } // namespace
849 } // namespace sparse
850 } // namespace tensorflow
851