| 88 | namespace { |
| 89 | |
| 90 | Status CopyVariable(int output_idx, OpKernelContext* ctx, const Tensor* t) { |
| 91 | Tensor* output; |
| 92 | Notification n; |
| 93 | Status status; |
| 94 | AllocatorAttributes attr; |
| 95 | if (t->dtype() == DT_VARIANT) { |
| 96 | attr.set_on_host(true); |
| 97 | } |
| 98 | TF_RETURN_IF_ERROR( |
| 99 | ctx->allocate_output(output_idx, t->shape(), &output, attr)); |
| 100 | if (t->dtype() == DT_VARIANT) { |
| 101 | output->flat<Variant>() = t->flat<Variant>(); |
| 102 | } else if (ctx->op_device_context() != nullptr) { |
| 103 | // TODO(apassos): remove the down_cast by just returning Device* from |
| 104 | // OpKernelContext |
| 105 | Device* device = static_cast<Device*>(ctx->device()); |
| 106 | ctx->op_device_context()->CopyTensorInSameDevice( |
| 107 | t, device, output, [&n, &status](const Status& s) { |
| 108 | status = s; |
| 109 | n.Notify(); |
| 110 | }); |
| 111 | n.WaitForNotification(); |
| 112 | return status; |
| 113 | } else { |
| 114 | switch (t->dtype()) { |
| 115 | #define HANDLER(type) \ |
| 116 | case DataTypeToEnum<type>::value: \ |
| 117 | output->flat<type>() = t->flat<type>(); \ |
| 118 | break; |
| 119 | TF_CALL_ALL_TYPES(HANDLER); |
| 120 | #undef HANDLER |
| 121 | default: |
| 122 | return errors::Internal("Unsupported dtype", t->dtype()); |
| 123 | } |
| 124 | } |
| 125 | return Status::OK(); |
| 126 | } |
| 127 | |
| 128 | } // namespace |
| 129 |
no test coverage detected