blob: 7f8b070556030d1c30956207e3c247a049286055 (
plain)
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
|
#include "tensorflow/core/kernels/initializable_lookup_table.h"
#include "tensorflow/core/lib/core/errors.h"
namespace tensorflow {
namespace lookup {
Status InitializableLookupTable::Find(const Tensor& keys, Tensor* values,
const Tensor& default_value) {
if (!is_initialized()) {
return errors::FailedPrecondition("Table not initialized.");
}
TF_RETURN_IF_ERROR(CheckFindArguments(keys, *values, default_value));
return DoFind(keys, values, default_value);
}
Status InitializableLookupTable::Initialize(InitTableIterator& iter) {
if (!iter.Valid()) {
return iter.status();
}
TF_RETURN_IF_ERROR(CheckKeyAndValueTensors(iter.keys(), iter.values()));
mutex_lock l(mu_);
if (is_initialized()) {
return errors::FailedPrecondition("Table already initialized.");
}
TF_RETURN_IF_ERROR(DoPrepare(iter.total_size()));
while (iter.Valid()) {
TF_RETURN_IF_ERROR(DoInsert(iter.keys(), iter.values()));
iter.Next();
}
if (!errors::IsOutOfRange(iter.status())) {
return iter.status();
}
is_initialized_ = true;
return Status::OK();
}
} // namespace lookup
} // namespace tensorflow
|