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