Searched refs:xla_tensor (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_device_context.cc | 125 XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); in CopyCPUTensorToDevice() local
126 CHECK(xla_tensor); in CopyCPUTensorToDevice()
134 TF_RET_CHECK(!xla_tensor->has_shaped_buffer()); in CopyCPUTensorToDevice()
136 xla_tensor->set_host_tensor(*cpu_tensor); in CopyCPUTensorToDevice()
138 xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_, in CopyCPUTensorToDevice()
152 << xla_tensor->shaped_buffer().ToString(); in CopyCPUTensorToDevice()
155 stream_->parent(), xla_tensor->shaped_buffer())) { in CopyCPUTensorToDevice()
162 host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer())); in CopyCPUTensorToDevice()
168 xla_tensor->ResetDefinitionEvent(std::move(event), in CopyCPUTensorToDevice()
231 XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); in CopyDeviceTensorToCPU() local
[all …]
|
| D | xla_tensor.cc | 25 XlaTensor* xla_tensor = in FromTensor() local
27 return xla_tensor; in FromTensor()
36 const XlaTensor* xla_tensor = FromTensor(&tensor); in DeviceMemoryFromTensor() local
37 if (xla_tensor) { in DeviceMemoryFromTensor()
38 CHECK(xla_tensor->has_shaped_buffer()); in DeviceMemoryFromTensor()
39 return xla_tensor->shaped_buffer().root_buffer(); in DeviceMemoryFromTensor()
|
| D | xla_launch_util.cc | 233 XlaTensor* xla_tensor = XlaTensor::FromTensor(t); in PopulateInputs() local
234 CHECK(xla_tensor); in PopulateInputs()
235 xla_tensor->WaitForDefinitionEventOnStream(stream); in PopulateInputs()
241 const XlaTensor* xla_tensor = XlaTensor::FromTensor(t); in PopulateInputs() local
242 CHECK(xla_tensor && xla_tensor->has_shaped_buffer()); in PopulateInputs()
243 arg_ptrs_[i] = const_cast<ShapedBuffer*>(&xla_tensor->shaped_buffer()); in PopulateInputs()
341 if (XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor)) { in PopulateOutputs() local
342 xla_tensor->set_host_tensor(const_tensor); in PopulateOutputs()
358 XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor); in PopulateOutputs() local
359 if (xla_tensor) { in PopulateOutputs()
[all …]
|
| D | xla_compile_on_demand_op.cc | 120 if (const XlaTensor* xla_tensor = XlaTensor::FromTensor(&device_tensor)) { in Compile() local
121 if (xla_tensor->has_host_tensor()) { in Compile()
126 constant_arguments[i] = xla_tensor->host_tensor(); in Compile()
|
| D | xla_device.cc | 118 const tensorflow::XlaTensor* xla_tensor = in DefaultPaddedShapeFn() local
120 if (xla_tensor == nullptr) { in DefaultPaddedShapeFn()
124 const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer(); in DefaultPaddedShapeFn()
|
| D | BUILD | 133 name = "xla_tensor",
134 srcs = ["xla_tensor.cc"],
135 hdrs = ["xla_tensor.h"],
168 ":xla_tensor",
261 ":xla_tensor",
|