/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/ |
H A D | dtensor_device_util.h | 146 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 D | dtensor_device.cc | 258 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 D | dtensor_device_util.cc | 56 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 D | BUILD | 104 "//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 D | parallel_device_lib_test.cc | 30 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 D | parallel_device.cc | 33 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 D | BUILD | 28 srcs = ["parallel_device.h"], 33 srcs = ["parallel_device.cc"], 55 name = "parallel_device", 118 ":parallel_device", 133 ":parallel_device", 155 ":parallel_device",
|
H A D | parallel_device_lib.h | 35 namespace parallel_device { 181 const ParallelDevice& parallel_device, 186 const ParallelDevice& parallel_device,
|
H A D | parallel_device_lib.cc | 32 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 D | parallel_device_testlib.h | 30 namespace parallel_device { 159 tensorflow::parallel_device::AllocateParallelDevice( in RegisterParallelDevice()
|
H A D | parallel_device.h | 24 namespace parallel_device {
|
H A D | parallel_device_remote_test.cc | 45 namespace parallel_device { namespace
|
H A D | parallel_device_testlib.cc | 32 namespace parallel_device { namespace
|
H A D | parallel_device_test.cc | 37 namespace parallel_device { namespace
|
/aosp_15_r20/external/tensorflow/tensorflow/python/distribute/parallel_device/ |
H A D | BUILD | 12 name = "parallel_device", 13 srcs = ["parallel_device.py"], 38 ":parallel_device",
|
H A D | parallel_device.py | 40 parallel_device = _all_parallel_devices.get(tensor.device, None) 41 if parallel_device is None: 43 return parallel_device.unpack(tensor)
|
H A D | parallel_device_test.py | 23 from tensorflow.python.distribute.parallel_device import parallel_device 116 self.device = parallel_device.ParallelDevice(components=[ 157 device = parallel_device.ParallelDevice(components=[
|
H A D | pywrap_parallel_device.cc | 55 tensorflow::parallel_device::AllocateParallelDevice( in PYBIND11_MODULE()
|
/aosp_15_r20/external/tensorflow/tensorflow/python/eager/ |
H A D | def_function.py | 72 from tensorflow.python.distribute.parallel_device import parallel_device 493 components = parallel_device.unpack(var_is_initialized[index])
|
H A D | BUILD | 949 "//tensorflow/python/distribute/parallel_device",
|
/aosp_15_r20/external/tensorflow/tensorflow/python/ |
H A D | BUILD | 4101 "//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/ |
D | files.db | [all...] |