Home | History | Annotate | Download | only in kernels
      1 /* Copyright 2016 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/framework/allocator.h"
     17 #include "tensorflow/core/framework/fake_input.h"
     18 #include "tensorflow/core/framework/node_def_builder.h"
     19 #include "tensorflow/core/framework/op_kernel.h"
     20 #include "tensorflow/core/framework/tensor.h"
     21 #include "tensorflow/core/framework/tensor_testutil.h"
     22 #include "tensorflow/core/framework/types.h"
     23 #include "tensorflow/core/kernels/ops_testutil.h"
     24 #include "tensorflow/core/platform/test.h"
     25 
     26 namespace tensorflow {
     27 
     28 namespace {
     29 
     30 class SparseReduceSumOpTest : public OpsTestBase {
     31  protected:
     32   template <typename T>
     33   void MakeOp() {
     34     DataType value_type = tensorflow::DataTypeToEnum<T>::value;
     35     TF_ASSERT_OK(NodeDefBuilder("sparse_reduce_sum", "SparseReduceSum")
     36                      .Input(FakeInput(DT_INT64))
     37                      .Input(FakeInput(value_type))
     38                      .Input(FakeInput(DT_INT64))
     39                      .Input(FakeInput(DT_INT32))
     40                      .Attr("T", value_type)
     41                      .Finalize(node_def()));
     42     TF_ASSERT_OK(InitOp());
     43   }
     44 };
     45 
     46 TEST_F(SparseReduceSumOpTest, SimpleReduce) {
     47   MakeOp<float>();
     48 
     49   // [    1]
     50   // [2    ]
     51   // [3   4]
     52 
     53   const auto indices_shape = TensorShape({4, 2});
     54   std::initializer_list<int64> in{0, 1, 1, 0, 2, 0, 2, 1};
     55   const gtl::ArraySlice<int64> indices(in);
     56   std::initializer_list<int64> sh{3, 2};
     57   const gtl::ArraySlice<int64> shape(sh);
     58 
     59   AddInputFromArray<int64>(indices_shape, indices);
     60   AddInputFromArray<float>(TensorShape({4}), {1, 2, 3, 4});
     61   AddInputFromArray<int64>(TensorShape({2}), shape);
     62   AddInputFromArray<int32>(TensorShape({1}), {0});  // reduction axes
     63 
     64   TF_ASSERT_OK(RunOpKernel());
     65 
     66   Tensor expected_by_0(allocator(), DT_FLOAT, TensorShape({2}));
     67   test::FillValues<float>(&expected_by_0, {5, 5});
     68   test::ExpectTensorEqual<float>(expected_by_0, *GetOutput(0));
     69 }
     70 
     71 class SparseReduceSumSparseOpTest : public OpsTestBase {
     72  protected:
     73   template <typename T>
     74   void MakeOp() {
     75     DataType value_type = tensorflow::DataTypeToEnum<T>::value;
     76     TF_ASSERT_OK(
     77         NodeDefBuilder("sparse_reduce_sum_sparse", "SparseReduceSumSparse")
     78             .Input(FakeInput(DT_INT64))
     79             .Input(FakeInput(value_type))
     80             .Input(FakeInput(DT_INT64))
     81             .Input(FakeInput(DT_INT32))
     82             .Attr("T", value_type)
     83             .Finalize(node_def()));
     84     TF_ASSERT_OK(InitOp());
     85   }
     86 };
     87 
     88 TEST_F(SparseReduceSumSparseOpTest, SimpleReduce) {
     89   MakeOp<float>();
     90 
     91   // [    2]
     92   // [2    ]
     93   // [3   4]
     94 
     95   const auto indices_shape = TensorShape({4, 2});
     96   std::initializer_list<int64> in{0, 1, 1, 0, 2, 0, 2, 1};
     97   const gtl::ArraySlice<int64> indices(in);
     98   std::initializer_list<int64> sh{3, 2};
     99   const gtl::ArraySlice<int64> shape(sh);
    100 
    101   AddInputFromArray<int64>(indices_shape, indices);
    102   AddInputFromArray<float>(TensorShape({4}), {2, 2, 3, 4});
    103   AddInputFromArray<int64>(TensorShape({2}), shape);
    104   AddInputFromArray<int32>(TensorShape({1}), {0});  // reduction axes
    105 
    106   TF_ASSERT_OK(RunOpKernel());
    107 
    108   const int expected_nnz = 2;
    109   Tensor expected_indices(allocator(), DT_INT64,
    110                           TensorShape({expected_nnz, 1}));
    111   test::FillValues<int64>(&expected_indices, {0, 1});
    112   test::ExpectTensorEqual<int64>(expected_indices, *GetOutput(0));
    113 
    114   Tensor expected_values(allocator(), DT_FLOAT, TensorShape({2}));
    115   test::FillValues<float>(&expected_values, {5, 6});
    116   test::ExpectTensorEqual<float>(expected_values, *GetOutput(1));
    117 
    118   Tensor expected_shape(allocator(), DT_INT64, {1});
    119   test::FillValues<int64>(&expected_shape, {2});
    120   test::ExpectTensorEqual<int64>(expected_shape, *GetOutput(2));
    121 }
    122 
    123 }  // namespace
    124 
    125 }  // namespace tensorflow
    126