• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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