| 42 | : OpKernel(context) {} |
| 43 | |
| 44 | void Compute(OpKernelContext* ctx) override { |
| 45 | const Tensor& val = ctx->input(0); |
| 46 | auto session_state = ctx->session_state(); |
| 47 | OP_REQUIRES(ctx, session_state != nullptr, |
| 48 | errors::FailedPrecondition( |
| 49 | "GetSessionHandle called on null session state")); |
| 50 | int64 id = session_state->GetNewId(); |
| 51 | TensorStore::TensorAndKey tk{val, id, requested_device()}; |
| 52 | OP_REQUIRES_OK(ctx, ctx->tensor_store()->AddTensor(name(), tk)); |
| 53 | |
| 54 | Tensor* handle = nullptr; |
| 55 | OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &handle)); |
| 56 | if (ctx->expected_output_dtype(0) == DT_RESOURCE) { |
| 57 | ResourceHandle resource_handle = MakeResourceHandle<Tensor>( |
| 58 | ctx, SessionState::kTensorHandleResourceTypeName, |
| 59 | tk.GetHandle(name())); |
| 60 | resource_handle.set_maybe_type_name( |
| 61 | SessionState::kTensorHandleResourceTypeName); |
| 62 | handle->scalar<ResourceHandle>()() = resource_handle; |
| 63 | } else { |
| 64 | // Legacy behavior in V1. |
| 65 | handle->flat<tstring>().setConstant(tk.GetHandle(name())); |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | TF_DISALLOW_COPY_AND_ASSIGN(GetSessionHandleOp); |
| 70 | }; |
nothing calls this directly
no test coverage detected