MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / CopyVariable

Function CopyVariable

tensorflow/core/kernels/resource_variable_ops.cc:90–126  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

88namespace {
89
90Status 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

Callers 1

ComputeMethod · 0.85

Calls 10

InternalFunction · 0.85
set_on_hostMethod · 0.80
allocate_outputMethod · 0.80
dtypeMethod · 0.45
shapeMethod · 0.45
op_device_contextMethod · 0.45
deviceMethod · 0.45
NotifyMethod · 0.45
WaitForNotificationMethod · 0.45

Tested by

no test coverage detected