1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
|
/* Copyright 2015 Google Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#define EIGEN_USE_THREADS
#include <string>
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/initializable_lookup_table.h"
#include "tensorflow/core/kernels/lookup_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/public/status.h"
#include "tensorflow/core/public/tensor.h"
namespace tensorflow {
namespace lookup {
// Iterator to initialize tables given 'keys' and 'values' tensors.
//
// The two tensors are returned in the first iteration. It doesn't loop
// over each element of the tensor since insertions in the lookup table can
// process batches.
class KeyValueTensorIterator
: public InitializableLookupTable::InitTableIterator {
public:
// keys and values are not owned by the iterator.
explicit KeyValueTensorIterator(const Tensor* keys, const Tensor* values)
: keys_(keys), values_(values), valid_(true), status_(Status::OK()) {
TensorShape key_shape = keys_->shape();
if (!key_shape.IsSameSize(values_->shape())) {
valid_ = false;
status_ = errors::InvalidArgument(
"keys and values should have the same dimension.",
key_shape.DebugString(), " vs ", values_->shape().DebugString());
}
if (key_shape.num_elements() == 0) {
valid_ = false;
status_ =
errors::InvalidArgument("keys and values cannot be empty tensors.");
}
}
bool Valid() const override { return valid_; }
void Next() override {
valid_ = false;
status_ = errors::OutOfRange("No more data.");
}
const Tensor& keys() const override { return *keys_; }
const Tensor& values() const override { return *values_; }
Status status() const override { return status_; }
int64 total_size() const {
return keys_ == nullptr ? -1 : keys_->NumElements();
}
private:
TF_DISALLOW_COPY_AND_ASSIGN(KeyValueTensorIterator);
const Tensor* keys_; // Doesn't own it.
const Tensor* values_; // Doesn't own it.
bool valid_; // true if the iterator points to an existing range.
Status status_;
};
} // namespace lookup
// Kernel to initialize a look table given a key and value tensors.
// After this operation, the table becomes read-only.
class InitializeTableOp : public OpKernel {
public:
explicit InitializeTableOp(OpKernelConstruction* context)
: OpKernel(context) {}
void Compute(OpKernelContext* ctx) override {
mutex_lock l(mu_);
lookup::InitializableLookupTable* table;
OP_REQUIRES_OK(ctx,
GetInitializableLookupTable("table_handle", ctx, &table));
core::ScopedUnref unref_me(table);
DataTypeVector expected_inputs = {DT_STRING_REF, table->key_dtype(),
table->value_dtype()};
DataTypeVector expected_outputs = {};
OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));
const Tensor& keys = ctx->input(1);
OP_REQUIRES(ctx, TensorShapeUtils::IsVector(keys.shape()),
errors::InvalidArgument("Keys must be a vector, but received ",
keys.shape().DebugString()));
const Tensor& values = ctx->input(2);
OP_REQUIRES(
ctx, TensorShapeUtils::IsVector(values.shape()),
errors::InvalidArgument("Values must be a vector, but received ",
values.shape().DebugString()));
OP_REQUIRES(ctx, keys.NumElements() == values.NumElements(),
errors::InvalidArgument(
"Keys and values must have the same size ",
keys.NumElements(), " vs ", values.NumElements()));
lookup::KeyValueTensorIterator iter(&keys, &values);
OP_REQUIRES_OK(ctx, table->Initialize(iter));
}
private:
mutex mu_;
};
REGISTER_KERNEL_BUILDER(Name("InitializeTable").Device(DEVICE_CPU),
InitializeTableOp);
} // namespace tensorflow
|
