| 858 | explicit LookupTableRemoveOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} |
| 859 | |
| 860 | void Compute(OpKernelContext* ctx) override { |
| 861 | lookup::LookupInterface* table; |
| 862 | OP_REQUIRES_OK(ctx, GetLookupTable("table_handle", ctx, &table)); |
| 863 | core::ScopedUnref unref_me(table); |
| 864 | |
| 865 | DataType expected_input_0 = |
| 866 | (ctx->input_dtype(0) == DT_RESOURCE) ? DT_RESOURCE : DT_STRING_REF; |
| 867 | DataTypeVector expected_inputs = {expected_input_0, table->key_dtype()}; |
| 868 | OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {})); |
| 869 | |
| 870 | const Tensor& key = ctx->input(1); |
| 871 | OP_REQUIRES_OK(ctx, table->CheckKeyTensorForRemove(key)); |
| 872 | |
| 873 | int64 memory_used_before = 0; |
| 874 | if (ctx->track_allocations()) { |
| 875 | memory_used_before = table->MemoryUsed(); |
| 876 | } |
| 877 | OP_REQUIRES_OK(ctx, table->Remove(ctx, key)); |
| 878 | if (ctx->track_allocations()) { |
| 879 | ctx->record_persistent_memory_allocation(table->MemoryUsed() - |
| 880 | memory_used_before); |
| 881 | } |
| 882 | } |
| 883 | }; |
| 884 | |
| 885 | REGISTER_KERNEL_BUILDER(Name("LookupTableRemoveV2").Device(DEVICE_CPU), |
nothing calls this directly
no test coverage detected