Home
last modified time | relevance | path

Searched refs:parallel_device (Results 1 – 22 of 22) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Ddtensor_device_util.h146 std::unique_ptr<parallel_device::ParallelDevice> parallel_device,
149 parallel_device_(std::move(parallel_device)), in mesh_config_()
160 parallel_device::ParallelTensor* DeviceIDs(TFE_Context* context,
162 const parallel_device::ParallelDevice& parallel_device() const { in parallel_device() function
189 std::unique_ptr<parallel_device::ParallelDevice> parallel_device_;
200 mutable std::unique_ptr<parallel_device::ParallelTensor> device_ids_tensor_;
220 std::unique_ptr<parallel_device::ParallelTensor> tensor,
292 virtual parallel_device::ParallelTensor* tensor() const { in tensor()
328 TensorWithLayout(std::unique_ptr<parallel_device::ParallelTensor> tensor,
340 std::unique_ptr<parallel_device::ParallelTensor> tensor_;
[all …]
H A Ddtensor_device.cc258 pair.second->parallel_device().AsyncWait(context, in AsyncWait()
353 const std::vector<parallel_device::ParallelTensor*>& parallel_inputs,
429 parallel_device::ParallelTensor* MeshWithParallelDevice::DeviceIDs( in DeviceIDs()
751 std::vector<parallel_device::TensorHandlePtr>& components, in VerifyPackTensorShapeAndDtype()
851 target_parallel_device->parallel_device().num_underlying_devices()) { in Pack()
860 std::vector<parallel_device::TensorHandlePtr> components; in Pack()
883 std::unique_ptr<parallel_device::ParallelTensor> parallel_tensor = in Pack()
884 parallel_device::ParallelTensor::FromTensorHandles( in Pack()
885 target_parallel_device->parallel_device(), std::move(components), in Pack()
1000 std::vector<parallel_device::TensorHandlePtr> indices_components; in SparsePack()
[all …]
H A Ddtensor_device_util.cc56 std::unique_ptr<parallel_device::ParallelTensor>
66 std::vector<parallel_device::TensorHandlePtr> components; in BroadcastTensorHandleToParallelTensor()
83 std::unique_ptr<parallel_device::ParallelTensor> parallel_tensor = in BroadcastTensorHandleToParallelTensor()
84 parallel_device::ParallelTensor::FromTensorHandles( in BroadcastTensorHandleToParallelTensor()
85 mesh.parallel_device(), std::move(components), status); in BroadcastTensorHandleToParallelTensor()
149 std::unique_ptr<parallel_device::ParallelTensor> parallel_tensor = in BroadcastResourceTensor()
307 std::unique_ptr<parallel_device::ParallelTensor> parallel_tensor = in Broadcast()
332 std::unique_ptr<parallel_device::ParallelTensor> tensor, in Wrap()
434 std::unique_ptr<parallel_device::ParallelTensor> indices_tensor, in Wrap()
435 std::unique_ptr<parallel_device::ParallelTensor> values_tensor, in Wrap()
[all …]
H A DBUILD104 "//tensorflow/c/eager/parallel_device:parallel_device_lib",
285 "//tensorflow/c/eager/parallel_device:parallel_device_lib",
/aosp_15_r20/external/tensorflow/tensorflow/c/eager/parallel_device/
H A Dparallel_device_lib_test.cc30 namespace parallel_device { namespace
55 ParallelDevice parallel_device(std::move(devices)); in TEST() local
64 parallel_device.Execute(context.get(), std::vector<ParallelTensor*>(), in TEST()
78 parallel_device.Execute(context.get(), handle_inputs, "ReadVariableOp", in TEST()
85 parallel_device.Execute(context.get(), std::vector<ParallelTensor*>(), in TEST()
111 ParallelDevice parallel_device(std::move(devices)); in TEST() local
120 parallel_device.StartExecute(context.get(), std::vector<ParallelTensor*>(), in TEST()
124 auto outputs = parallel_device.Join( in TEST()
153 ParallelDevice parallel_device(devices); in TEST() local
189 parallel_device.ScalarsFromSequence<float>({1.0, 2.0}, context.get(), in TEST()
[all …]
H A Dparallel_device.cc33 namespace parallel_device { namespace
56 std::unique_ptr<ParallelDevice> parallel_device) in NamedParallelDevice() argument
57 : device_name_(name), parallel_device_(std::move(parallel_device)) {} in NamedParallelDevice()
67 const ParallelDevice& parallel_device, in ExecuteWithSpecialOps() argument
78 if (inputs.size() != parallel_device.num_underlying_devices()) { in ExecuteWithSpecialOps()
81 parallel_device.num_underlying_devices(), in ExecuteWithSpecialOps()
105 parallel_device, std::move(components), status)); in ExecuteWithSpecialOps()
116 if (expected_outputs != parallel_device.num_underlying_devices()) { in ExecuteWithSpecialOps()
119 parallel_device.num_underlying_devices(), in ExecuteWithSpecialOps()
163 parallel_device.CopyToParallelDevice( in ExecuteWithSpecialOps()
[all …]
H A DBUILD28 srcs = ["parallel_device.h"],
33 srcs = ["parallel_device.cc"],
55 name = "parallel_device",
118 ":parallel_device",
133 ":parallel_device",
155 ":parallel_device",
H A Dparallel_device_lib.h35 namespace parallel_device {
181 const ParallelDevice& parallel_device,
186 const ParallelDevice& parallel_device,
H A Dparallel_device_lib.cc32 namespace parallel_device { namespace
505 const ParallelDevice& parallel_device, in FromTensorHandles() argument
529 new ParallelTensor(parallel_device, std::move(components), shape, dtype)); in FromTensorHandles()
533 const ParallelDevice& parallel_device, in FromTensorHandles() argument
556 new ParallelTensor(parallel_device, std::move(components), dtype)); in FromTensorHandles()
H A Dparallel_device_testlib.h30 namespace parallel_device {
159 tensorflow::parallel_device::AllocateParallelDevice( in RegisterParallelDevice()
H A Dparallel_device.h24 namespace parallel_device {
H A Dparallel_device_remote_test.cc45 namespace parallel_device { namespace
H A Dparallel_device_testlib.cc32 namespace parallel_device { namespace
H A Dparallel_device_test.cc37 namespace parallel_device { namespace
/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/parallel_device/
H A DBUILD12 name = "parallel_device",
13 srcs = ["parallel_device.py"],
38 ":parallel_device",
H A Dparallel_device.py40 parallel_device = _all_parallel_devices.get(tensor.device, None)
41 if parallel_device is None:
43 return parallel_device.unpack(tensor)
H A Dparallel_device_test.py23 from tensorflow.python.distribute.parallel_device import parallel_device
116 self.device = parallel_device.ParallelDevice(components=[
157 device = parallel_device.ParallelDevice(components=[
H A Dpywrap_parallel_device.cc55 tensorflow::parallel_device::AllocateParallelDevice( in PYBIND11_MODULE()
/aosp_15_r20/external/tensorflow/tensorflow/python/eager/
H A Ddef_function.py72 from tensorflow.python.distribute.parallel_device import parallel_device
493 components = parallel_device.unpack(var_is_initialized[index])
H A DBUILD949 "//tensorflow/python/distribute/parallel_device",
/aosp_15_r20/external/tensorflow/tensorflow/python/
H A DBUILD4101 "//tensorflow/c/eager/parallel_device:headers",
4102 "//tensorflow/c/eager/parallel_device:sources",
4103 "//tensorflow/python/distribute/parallel_device:pywrap_parallel_device.cc",
4106 visibility = ["//tensorflow/python/distribute/parallel_device:__pkg__"],
/aosp_15_r20/out/.module_paths/
Dfiles.db[all...]