| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | hlo_runner.cc | 196 DeviceAssignment* device_assignment) { in ExecuteReplicated()
211 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl()
343 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated()
395 DeviceAssignment* device_assignment) { in ExecuteReplicated()
473 int64_t device, se::Stream* stream, DeviceAssignment* device_assignment, in GetServiceRunOptionsForDevice()
|
| H A D | collective_ops_utils.cc | 130 GetParticipatingDevicesGroups(const DeviceAssignment& device_assignment, in GetParticipatingDevicesGroups()
239 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices()
|
| H A D | collective_ops_utils_test.cc | 120 xla::Array2D<int> device_assignment; member
376 DeviceAssignment device_assignment(num_replicas, num_partitions); in TEST_P() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_compile_op_support.cc | 61 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
140 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
315 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules()
428 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext() local
|
| H A D | tpu_executable_info.proto | 63 xla.DeviceAssignmentProto device_assignment = 6; field
|
| H A D | tpu_functional_ops.h | 147 std::vector<int> device_assignment; member
|
| H A D | tpu_reshard_variables_op.cc | 218 std::unique_ptr<xla::DeviceAssignment> device_assignment; in DoTpuExecute() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/ |
| H A D | executable_run_options.cc | 96 const DeviceAssignment* device_assignment) { in set_device_assignment()
101 const DeviceAssignment* ExecutableRunOptions::device_assignment() const { in device_assignment() function in xla::ExecutableRunOptions
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| H A D | xrt_compile_ops.cc | 56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment()
171 xla::DeviceAssignment device_assignment(num_replicas, in Compile() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | all_reduce_blueconnect.cc | 61 const DeviceAssignment& device_assignment, size_t num_devices_per_host) { in TryDecomposeReplicaGroup()
112 const DeviceAssignment& device_assignment = in TryDecomposeReplicaGroups() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| H A D | tpu_client.cc | 522 DeviceAssignment device_assignment, std::shared_ptr<PyTpuClient> client, in PyTpuExecutable()
592 xla::DeviceAssignmentProto device_assignment; in ExecuteHelper() local
820 std::optional<xla::DeviceAssignment> device_assignment; in Compile() local
|
| H A D | libtpu_client.c | 117 DeviceAssignment device_assignment = {NULL, 0}; in main() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| H A D | pjrt_stream_executor_client.cc | 1764 std::shared_ptr<DeviceAssignment> device_assignment, in PjRtStreamExecutorExecutable()
1904 std::shared_ptr<DeviceAssignment> device_assignment, in EnqueueExecution()
2096 std::shared_ptr<DeviceAssignment> device_assignment; in ExecuteHelper() local
2391 std::shared_ptr<DeviceAssignment>& device_assignment = in GetExecutableExtras() local
2454 std::shared_ptr<DeviceAssignment>& device_assignment = in Compile() local
|
| H A D | pjrt_stream_executor_client.h | 332 std::shared_ptr<DeviceAssignment> device_assignment; member
742 const DeviceAssignment& device_assignment() const override { in device_assignment() function
|
| H A D | compile_options.proto | 50 xla.DeviceAssignmentProto device_assignment = 9; field
|
| H A D | gpu_multistream_test.cc | 54 DeviceAssignment device_assignment(1, 1); in TEST() local
|
| H A D | tfrt_cpu_pjrt_client.cc | 339 std::shared_ptr<DeviceAssignment> device_assignment; in Compile() local
1159 std::shared_ptr<DeviceAssignment> device_assignment, in TfrtCpuExecutable()
1311 std::shared_ptr<DeviceAssignment> device_assignment; in ExecuteHelper() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
| H A D | tpu_feed.py | 647 device_assignment=None, argument
772 device_assignment, argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | tpu_strategy.py | 682 device_assignment=None): argument
743 device_assignment=None): argument
846 device_assignment=None, argument
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | executable_build_options.cc | 163 const DeviceAssignment& device_assignment) { in set_device_assignment()
|
| H A D | executable_build_options.h | 123 const DeviceAssignment& device_assignment() const { in device_assignment() function
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| H A D | update_tpu_metadata.cc | 124 xla::DeviceAssignmentProto device_assignment; in UpdateTPUCompileMetadata() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/ |
| H A D | hlo_test_base.cc | 263 int64_t num_replicas, DeviceAssignment* device_assignment, in ExecuteReplicated()
281 DeviceAssignment* device_assignment) { in ExecuteReplicated()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/jit/ |
| H A D | xla_compile_on_demand_op.cc | 83 xla::DeviceAssignment device_assignment; in Run() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/protobuf/tpu/ |
| H A D | compile_metadata.proto | 90 xla.DeviceAssignmentProto device_assignment = 8; field
|