Home
last modified time | relevance | path

Searched refs:device_tensor (Results 1 – 20 of 20) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/
H A Dxla_device_context.cc116 Tensor* device_tensor, in CopyCPUTensorToDevice() argument
128 << reinterpret_cast<const void*>(device_tensor->tensor_data().data()) in CopyCPUTensorToDevice()
131 << device_tensor->shape().DebugString(); in CopyCPUTensorToDevice()
133 XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); in CopyCPUTensorToDevice()
138 device_tensor->shape(), device_tensor->dtype(), std::nullopt); in CopyCPUTensorToDevice()
142 device_tensor->shape(), device_tensor->dtype(), in CopyCPUTensorToDevice()
149 xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_, in CopyCPUTensorToDevice()
209 void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU() argument
213 if (device_tensor->NumElements() == 0) { in CopyDeviceTensorToCPU()
219 << reinterpret_cast<const void*>(device_tensor->tensor_data().data()) in CopyDeviceTensorToCPU()
[all …]
H A Dxla_device_context.h66 Tensor* device_tensor, StatusCallback done,
68 void CopyDeviceTensorToCPU(const Tensor* device_tensor,
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/
H A Dvirtual_device.cc28 Tensor* device_tensor, StatusCallback done,
30 void CopyDeviceTensorToCPU(const Tensor* device_tensor,
40 Tensor* device_tensor, in CopyCPUTensorToDevice() argument
43 *device_tensor = *cpu_tensor; in CopyCPUTensorToDevice()
47 void VirtualDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU() argument
52 *cpu_tensor = *device_tensor; in CopyDeviceTensorToCPU()
/aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/pluggable_device/
H A Dpluggable_device_util.cc176 const Tensor* device_tensor, Tensor* cpu_tensor, StatusCallback done) { in CopyPluggableDeviceTensorToCPU() argument
180 Status s = PrepareCopy(device, device_context, *device_tensor, cpu_tensor, in CopyPluggableDeviceTensorToCPU()
198 const int64_t total_bytes = device_tensor->TotalBytes(); in CopyPluggableDeviceTensorToCPU()
200 void* src_ptr = GetBase(device_tensor); in CopyPluggableDeviceTensorToCPU()
208 TensorReference input_ref(*device_tensor); in CopyPluggableDeviceTensorToCPU()
223 Device* device, Tensor* device_tensor, StatusCallback done, in CopyCPUTensorToPluggableDevice() argument
228 Status s = PrepareCopy(device, device_context, *cpu_tensor, device_tensor, in CopyCPUTensorToPluggableDevice()
251 void* dst_ptr = GetBase(device_tensor); in CopyCPUTensorToPluggableDevice()
H A Dpluggable_device_context.cc28 const Tensor* cpu_tensor, Device* device, Tensor* device_tensor, in CopyCPUTensorToDevice() argument
31 cpu_tensor, this, device, device_tensor, done, sync_dst_compute); in CopyCPUTensorToDevice()
34 void PluggableDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU() argument
40 device, this, device_tensor, cpu_tensor, done); in CopyDeviceTensorToCPU()
H A Dpluggable_device_util.h38 const Tensor* device_tensor, Tensor* cpu_tensor, StatusCallback done);
51 Device* device, Tensor* device_tensor, StatusCallback done,
H A Dpluggable_device_context.h53 Tensor* device_tensor, StatusCallback done,
56 void CopyDeviceTensorToCPU(const Tensor* device_tensor,
/aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/gpu/
H A Dgpu_util_platform_specific.cc28 Tensor* device_tensor, in CopyCPUTensorToDevice() argument
31 GPUUtil::CopyCPUTensorToGPU(cpu_tensor, this, device, device_tensor, done, in CopyCPUTensorToDevice()
35 void GPUDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU() argument
39 GPUUtil::CopyGPUTensorToCPU(device, this, device_tensor, cpu_tensor, done); in CopyDeviceTensorToCPU()
/aosp_15_r20/external/tensorflow/tensorflow/core/framework/
H A Ddevice_base.cc37 Status DeviceContext::CopyDeviceTensorToCPUSync(const Tensor* device_tensor, in CopyDeviceTensorToCPUSync() argument
43 CopyDeviceTensorToCPU(device_tensor, tensor_name, device, cpu_tensor, in CopyDeviceTensorToCPUSync()
54 Tensor* device_tensor) const { in CopyCPUTensorToDeviceSync()
57 CopyCPUTensorToDevice(cpu_tensor, device, device_tensor, in CopyCPUTensorToDeviceSync()
H A Ddevice_base.h80 Tensor* device_tensor, StatusCallback done,
87 Tensor* device_tensor) const;
99 virtual void CopyDeviceTensorToCPU(const Tensor* device_tensor, in CopyDeviceTensorToCPU() argument
106 Status CopyDeviceTensorToCPUSync(const Tensor* device_tensor,
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dstack.cc306 Tensor* device_tensor = in ComputeAsync() local
309 cpu_tensor, device, device_tensor, in ComputeAsync()
310 [device_tensor, ctx, done](const Status& s) { in ComputeAsync()
313 ctx->set_output(0, *device_tensor); in ComputeAsync()
316 delete device_tensor; in ComputeAsync()
H A Dresource_variable_ops.cc941 Status CopyTensorToHost(OpKernelContext* c, const Tensor& device_tensor, in CopyTensorToHost() argument
948 device_tensor.dtype(), device_tensor.shape(), host_tensor, alloc_attr)); in CopyTensorToHost()
950 const_cast<Tensor&>(device_tensor).flat<T>().data(), in CopyTensorToHost()
951 device_tensor.flat<T>().size() * sizeof(T)); in CopyTensorToHost()
953 device_tensor.NumElements() * sizeof(T)); in CopyTensorToHost()
/aosp_15_r20/external/tensorflow/tensorflow/core/common_runtime/
H A Dfunction_testlib.cc33 Tensor* device_tensor = nullptr; in Compute() local
35 &device_tensor)); in Compute()
36 device_tensor->scalar<tstring>()() = in Compute()
H A Dgpu_device_context.h68 Tensor* device_tensor, StatusCallback done,
71 void CopyDeviceTensorToCPU(const Tensor* device_tensor, StringPiece edge_name,
H A Dprocess_function_library_runtime_test.cc176 Tensor GPUToCPU(const Tensor& device_tensor) { in GPUToCPU() argument
183 Tensor cpu_tensor(device_tensor.dtype(), device_tensor.shape()); in GPUToCPU()
185 ->CopyDeviceTensorToCPUSync(&device_tensor, "", gpu_device_, in GPUToCPU()
201 Tensor device_tensor(gpu_device_->GetAllocator({}), cpu_tensor.dtype(), in CPUToGPU() local
205 &device_tensor) in CPUToGPU()
207 return device_tensor; in CPUToGPU()
/aosp_15_r20/external/pytorch/torch/csrc/
H A Dserialization.cpp241 auto device_tensor = at::from_blob( in THPStorage_writeFileRaw() local
248 cpu_tensor = device_tensor.to(at::kCPU); in THPStorage_writeFileRaw()
390 auto device_tensor = at::from_blob( in THPStorage_readFileRaw() local
397 device_tensor.copy_(cpu_tensor); in THPStorage_readFileRaw()
/aosp_15_r20/external/pytorch/torch/csrc/lazy/ts_backend/
H A Dts_backend_impl.cpp102 auto device_tensor = at::full(tensor.sizes(), tensor.item(), options); in MakeComputationDataFromTensor() local
103 return std::make_shared<TSData>(device_tensor, shape, device); in MakeComputationDataFromTensor()
/aosp_15_r20/external/pytorch/test/
H A Dtest_reductions.py2013 device_tensor = cpu_tensor.to(device)
2014 self.assertEqual(device_tensor.var(), cpu_tensor.var())
2015 self.assertEqual(device_tensor.var(1), cpu_tensor.var(1))
2016 self.assertEqual(device_tensor.var(2), cpu_tensor.var(2))
2017 self.assertEqual(device_tensor.std(), cpu_tensor.std())
2018 self.assertEqual(device_tensor.std(1), cpu_tensor.std(1))
2019 self.assertEqual(device_tensor.var(2), cpu_tensor.var(2))
2022 device_tensor = cpu_tensor.to(device)
2023 self.assertEqual(device_tensor.var(), cpu_tensor.var())
2030 device_tensor = cpu_tensor.to(device)
[all …]
H A Dtest_tensor_creation_ops.py2896 device_tensor = torch.arange(0, 10, dtype=dtype, device=device)
2897 self.assertEqual(cpu_tensor, device_tensor)
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2tensorrt/kernels/
H A Dtrt_engine_op.cc671 const auto device_tensor = &calib_ctx->device_tensors_.at(i); in ExecuteCalibration() local
672 CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); in ExecuteCalibration()