Home
last modified time | relevance | path

Searched defs:device_assignment (Results 1 – 25 of 43) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dcollective_ops_utils_test.cc53 DeviceAssignment device_assignment(/*replica_count=*/3, in TEST() local
69 DeviceAssignment device_assignment(/*replica_count=*/4, in TEST() local
92 DeviceAssignment device_assignment(/*replica_count=*/2, in TEST() local
Dhlo_runner.cc185 DeviceAssignment* device_assignment) { in ExecuteReplicated()
200 DeviceAssignment* device_assignment) { in ExecuteReplicatedImpl()
321 DeviceAssignment* device_assignment, ExecutionProfile* profile) { in ExecuteReplicated()
443 int64 device, se::Stream* stream, DeviceAssignment* device_assignment, in GetServiceRunOptionsForDevice()
Dcollective_ops_utils.cc79 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices()
Dcompiler.h103 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
Dhlo_module_config.h185 void set_static_device_assignment(const DeviceAssignment& device_assignment) { in set_static_device_assignment()
Dservice.cc362 DeviceAssignment device_assignment(options_.number_of_replicas(), in ExecuteParallelAndRegisterResult() local
485 DeviceAssignment device_assignment(options_.number_of_replicas(), in ExecuteAndRegisterResult() local
/external/tensorflow/tensorflow/compiler/xla/
Dexecutable_run_options.cc96 const DeviceAssignment* device_assignment) { in set_device_assignment()
101 const DeviceAssignment* ExecutableRunOptions::device_assignment() const { in device_assignment() function in xla::ExecutableRunOptions
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.cc59 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
138 absl::optional<const DeviceAssignment> device_assignment, int replica_count, in CreateModuleConfig()
310 const absl::optional<xla::DeviceAssignment>& device_assignment, in CreateHloModules()
422 const DeviceAssignment& device_assignment = in CompileOpMetadataFromContext() local
Dtpu_executable_info.proto63 xla.DeviceAssignmentProto device_assignment = 6; field
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_compile_ops.cc56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment()
170 xla::DeviceAssignment device_assignment(num_replicas, in Compile() local
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dgpu_multistream_test.cc53 DeviceAssignment device_assignment(1, 1); in TEST() local
Dutils.cc118 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions()
/external/tensorflow/tensorflow/compiler/xla/client/
Dexecutable_build_options.cc93 const DeviceAssignment& device_assignment) { in set_device_assignment()
Dexecutable_build_options.h108 const DeviceAssignment& device_assignment() const { in device_assignment() function
/external/tensorflow/tensorflow/python/tpu/
Dtpu_feed.py647 device_assignment=None, argument
771 device_assignment, argument
Ddevice_assignment.py337 def device_assignment( function
/external/tensorflow/tensorflow/core/protobuf/tpu/
Dcompile_metadata.proto86 xla.DeviceAssignmentProto device_assignment = 8; field
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dlibtpu_client.c117 DeviceAssignment device_assignment = {NULL, 0}; in main() local
Dtpu_client.cc511 DeviceAssignment device_assignment, std::shared_ptr<PyTpuClient> client, in PyTpuExecutable()
581 xla::DeviceAssignmentProto device_assignment; in ExecuteHelper() local
810 absl::optional<xla::DeviceAssignment> device_assignment; in Compile() local
Dtpu_client.h322 const DeviceAssignment& device_assignment() const { in device_assignment() function
/external/tensorflow/tensorflow/python/distribute/
Dtpu_strategy.py627 device_assignment=None): argument
687 device_assignment=None): argument
790 device_assignment=None): argument
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/
Dtpu_service.proto94 optional xla.DeviceAssignmentProto device_assignment = 4; field
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt.proto50 DeviceAssignment device_assignment = 6; field
/external/tensorflow/tensorflow/compiler/xla/python/
Doutfeed_receiver.cc402 DeviceAssignment device_assignment(1, 1); in SendShutdownOutfeedHeader() local
/external/tensorflow/tensorflow/core/tpu/
Dtpu_execute.cc418 TpuNodeContext* node_context, xla::DeviceAssignment* device_assignment, in TPUExecute()

12