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