| 131 | } |
| 132 | |
| 133 | Status PartitionedCallOp::FillOutputDevices( |
| 134 | const FunctionLibraryRuntime& lib, const Device& cpu_device, |
| 135 | AttrSlice attrs, FunctionLibraryRuntime::InstantiateOptions* opts) { |
| 136 | const FunctionLibraryDefinition* flib = lib.GetFunctionLibraryDefinition(); |
| 137 | const FunctionDef* fdef = flib->Find(func_->name()); |
| 138 | if (fdef == nullptr) { |
| 139 | return errors::NotFound("Failed for find definition for function \"", |
| 140 | func_->name(), "\""); |
| 141 | } |
| 142 | |
| 143 | bool is_type_list; |
| 144 | for (const OpDef::ArgDef& ret_def : fdef->signature().output_arg()) { |
| 145 | DataTypeVector dtypes; |
| 146 | TF_RETURN_IF_ERROR(ArgNumType(attrs, ret_def, &is_type_list, &dtypes)); |
| 147 | for (DataType dtype : dtypes) { |
| 148 | if (dtype == DT_RESOURCE) { |
| 149 | // Resource memory type is HOST_MEMORY, however the actual resource |
| 150 | // might be allocated on a device. We leave output device for resource |
| 151 | // outputs empty, and rely on a Placer and colocation constraints to |
| 152 | // infer correct placement for the function output. |
| 153 | opts->output_devices.push_back(""); |
| 154 | } else if (MTypeFromDType(dtype) == HOST_MEMORY) { |
| 155 | opts->output_devices.push_back(cpu_device.name()); |
| 156 | } else { |
| 157 | opts->output_devices.push_back(opts->target); |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | return Status::OK(); |
| 162 | } |
| 163 | |
| 164 | Status PartitionedCallOp::Instantiate(FunctionLibraryRuntime* lib, |
| 165 | OpKernelContext* ctx, |
nothing calls this directly
no test coverage detected