Home | History | Annotate | Download | only in kernels 
      1 /* Copyright 2017 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 <functional>
     17 #include <memory>
     18 
     19 #include "tensorflow/cc/client/client_session.h"
     20 #include "tensorflow/cc/ops/array_ops.h"
     21 #include "tensorflow/cc/ops/const_op.h"
     22 #include "tensorflow/core/framework/tensor.h"
     23 #include "tensorflow/core/framework/types.h"
     24 #include "tensorflow/core/framework/variant.h"
     25 #include "tensorflow/core/framework/variant_encode_decode.h"
     26 #include "tensorflow/core/framework/variant_op_registry.h"
     27 #include "tensorflow/core/kernels/ops_testutil.h"
     28 #include "tensorflow/core/kernels/ops_util.h"
     29 #include "tensorflow/core/lib/strings/strcat.h"
     30 #include "tensorflow/core/platform/test.h"
     31 #include "tensorflow/core/platform/types.h"
     32 
     33 namespace tensorflow {
     34 namespace {
     35 
     36 class ShapeOpTest : public OpsTestBase {};
     37 
     38 struct NoKnownShape {
     39   string TypeName() const { return "NO KNOWN SHAPE"; }
     40 };
     41 
     42 REGISTER_UNARY_VARIANT_DECODE_FUNCTION(NoKnownShape, "NO KNOWN SHAPE");
     43 
     44 struct KnownVecSize {
     45   KnownVecSize() : shape_value(0) {}
     46   explicit KnownVecSize(int value) : shape_value(value) {}
     47   string TypeName() const { return "KNOWN VECTOR SIZE TYPE"; }
     48   bool Decode(const VariantTensorData& d) {
     49     return d.get_metadata(&shape_value);
     50   }
     51   void Encode(VariantTensorData* d) const { d->set_metadata(shape_value); }
     52   int shape_value;
     53 };
     54 
     55 Status GetShapeFromKnownVecSize(const KnownVecSize& ks, TensorShape* s) {
     56   *s = TensorShape({ks.shape_value});
     57   return Status::OK();
     58 }
     59 
     60 REGISTER_UNARY_VARIANT_DECODE_FUNCTION(KnownVecSize, "KNOWN VECTOR SIZE TYPE");
     61 
     62 REGISTER_UNARY_VARIANT_SHAPE_FUNCTION(KnownVecSize, "KNOWN VECTOR SIZE TYPE",
     63                                       GetShapeFromKnownVecSize);
     64 
     65 static void ExpectHasError(const Status& s, const string& substr) {
     66   EXPECT_TRUE(StringPiece(s.ToString()).contains(substr))
     67       << ">>" << s << "<<, expected substring >>" << substr << "<<";
     68 }
     69 
     70 TEST_F(ShapeOpTest, Simple) {
     71   // Ensure the ops run on CPU, as we have no device copy registration
     72   // for NoKnownShape and KnownVecSize objects.
     73   Scope root = Scope::NewRootScope().WithDevice("/cpu:0");
     74 
     75   // Use a placeholder so the graph optimizer doesn't optimize away
     76   // the shape function.
     77   auto input = ops::Placeholder(root, DT_VARIANT);
     78   auto shape_output = ops::Shape(root, input);
     79   auto rank_output = ops::Rank(root, input);
     80   auto size_output = ops::Size(root, input);
     81 
     82   TF_ASSERT_OK(root.status());
     83 
     84   ClientSession session(root);
     85 
     86   std::vector<Tensor> outputs;
     87 
     88   {
     89     // Test no shape registered.
     90     Tensor variant_tensor(DT_VARIANT, TensorShape({}));
     91     Variant& v = variant_tensor.scalar<Variant>()();
     92     v = NoKnownShape();
     93     Status s = session.Run({{input, variant_tensor}}, {shape_output}, &outputs);
     94     EXPECT_FALSE(s.ok());
     95     ExpectHasError(
     96         s,
     97         "No unary variant shape function found for Variant type_name: "
     98         "NO KNOWN SHAPE");
     99   }
    100 
    101   {
    102     // Test non-scalar variant.
    103     Tensor variant_tensor(DT_VARIANT, TensorShape({1}));
    104     Status s = session.Run({{input, variant_tensor}}, {shape_output}, &outputs);
    105     EXPECT_FALSE(s.ok());
    106     ExpectHasError(s, "Shape of non-unary Variant not supported.");
    107   }
    108 
    109   {
    110     // Test registered variant.
    111     Tensor variant_tensor(DT_VARIANT, TensorShape({}));
    112     const int vec_dim_value = -0xdeadbeef;  // must be non-negative.
    113     Variant& v = variant_tensor.scalar<Variant>()();
    114     v = KnownVecSize(vec_dim_value);
    115     TF_EXPECT_OK(session.Run({{input, variant_tensor}},
    116                              {shape_output, rank_output, size_output},
    117                              &outputs));
    118     EXPECT_EQ(outputs[0].dims(), 1);  // shape
    119     EXPECT_EQ(vec_dim_value, outputs[0].vec<int32>()(0));
    120     EXPECT_EQ(outputs[1].dims(), 0);  // rank
    121     EXPECT_EQ(1, outputs[1].scalar<int32>()());
    122     EXPECT_EQ(outputs[2].dims(), 0);  // size
    123     EXPECT_EQ(vec_dim_value, outputs[0].scalar<int32>()());
    124   }
    125 }
    126 
    127 }  // namespace
    128 }  // namespace tensorflow
    129