| 37 | // respectively. |
| 38 | template <class Container, class key_dtype, class value_dtype> |
| 39 | class LookupTableOp : public OpKernel { |
| 40 | public: |
| 41 | // ctx is not owned by this class. |
| 42 | explicit LookupTableOp(OpKernelConstruction* ctx) |
| 43 | : OpKernel(ctx), table_handle_set_(false) { |
| 44 | OP_REQUIRES_OK(ctx, ctx->allocate_persistent(tensorflow::DT_STRING, |
| 45 | tensorflow::TensorShape({2}), |
| 46 | &table_handle_, nullptr)); |
| 47 | OP_REQUIRES_OK( |
| 48 | ctx, ctx->GetAttr("use_node_name_sharing", &use_node_name_sharing_)); |
| 49 | } |
| 50 | |
| 51 | // ctx is not owned by this function. |
| 52 | void Compute(OpKernelContext* ctx) override { |
| 53 | mutex_lock l(mu_); |
| 54 | |
| 55 | if (!table_handle_set_) { |
| 56 | OP_REQUIRES_OK(ctx, cinfo_.Init(ctx->resource_manager(), def(), |
| 57 | use_node_name_sharing_)); |
| 58 | } |
| 59 | |
| 60 | auto creator = |
| 61 | [ctx, this](lookup::LookupInterface** ret) |
| 62 | EXCLUSIVE_LOCKS_REQUIRED(mu_) { |
| 63 | lookup::LookupInterface* container = new Container(ctx, this); |
| 64 | if (!ctx->status().ok()) { |
| 65 | container->Unref(); |
| 66 | return ctx->status(); |
| 67 | } |
| 68 | if (ctx->track_allocations()) { |
| 69 | ctx->record_persistent_memory_allocation( |
| 70 | container->MemoryUsed() + table_handle_.AllocatedBytes()); |
| 71 | } |
| 72 | *ret = container; |
| 73 | return Status::OK(); |
| 74 | }; |
| 75 | |
| 76 | lookup::LookupInterface* table = nullptr; |
| 77 | OP_REQUIRES_OK(ctx, |
| 78 | cinfo_.resource_manager() |
| 79 | ->template LookupOrCreate<lookup::LookupInterface>( |
| 80 | cinfo_.container(), cinfo_.name(), &table, creator)); |
| 81 | core::ScopedUnref unref_me(table); |
| 82 | |
| 83 | OP_REQUIRES_OK(ctx, lookup::CheckTableDataTypes( |
| 84 | *table, DataTypeToEnum<key_dtype>::v(), |
| 85 | DataTypeToEnum<value_dtype>::v(), cinfo_.name())); |
| 86 | |
| 87 | if (ctx->expected_output_dtype(0) == DT_RESOURCE) { |
| 88 | Tensor* handle; |
| 89 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &handle)); |
| 90 | handle->scalar<ResourceHandle>()() = |
| 91 | MakeResourceHandle<lookup::LookupInterface>(ctx, cinfo_.container(), |
| 92 | cinfo_.name()); |
| 93 | } else { |
| 94 | if (!table_handle_set_) { |
| 95 | auto h = table_handle_.AccessTensor(ctx)->template flat<tstring>(); |
| 96 | h(0) = cinfo_.container(); |
nothing calls this directly
no test coverage detected