Home
last modified time | relevance | path

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

123

/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/compiler/xla/service/
Dhlo_runner.cc196 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()
Dcollective_ops_utils.cc130 GetParticipatingDevicesGroups(const DeviceAssignment& device_assignment, in GetParticipatingDevicesGroups()
239 GlobalDeviceId device_id, const DeviceAssignment& device_assignment, in GetParticipatingDevices()
Dcollective_ops_utils_test.cc120 xla::Array2D<int> device_assignment; member
376 DeviceAssignment device_assignment(num_replicas, num_partitions); in TEST_P() local
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.cc61 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
Dtpu_executable_info.proto63 xla.DeviceAssignmentProto device_assignment = 6; field
Dtpu_functional_ops.h147 std::vector<int> device_assignment; member
Dtpu_reshard_variables_op.cc218 std::unique_ptr<xla::DeviceAssignment> device_assignment; in DoTpuExecute() local
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_compile_ops.cc56 int num_cores_per_replica, xla::DeviceAssignment* device_assignment) { in GenerateXlaDeviceAssignment()
171 xla::DeviceAssignment device_assignment(num_replicas, in Compile() local
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dcompile_options.proto50 xla.DeviceAssignmentProto device_assignment = 9; field
Dgpu_multistream_test.cc54 DeviceAssignment device_assignment(1, 1); in TEST() local
Dpjrt_stream_executor_client.h332 std::shared_ptr<DeviceAssignment> device_assignment; member
742 const DeviceAssignment& device_assignment() const override { in device_assignment() function
Dutils.cc120 std::shared_ptr<DeviceAssignment>* device_assignment) { in ParseDeviceAssignmentCompileOptions()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dall_reduce_blueconnect.cc61 const DeviceAssignment& device_assignment, size_t num_devices_per_host) { in TryDecomposeReplicaGroup()
112 const DeviceAssignment& device_assignment = in TryDecomposeReplicaGroups() local
Dall_reduce_blueconnect_test.cc38 DeviceAssignment device_assignment(replica_count, /*computation_count=*/1); in SetModuleConfig() local
/external/tensorflow/tensorflow/python/tpu/
Dtpu_feed.py647 device_assignment=None, argument
772 device_assignment, argument
Ddevice_assignment.py332 def device_assignment( function
/external/tensorflow/tensorflow/compiler/xla/client/
Dexecutable_build_options.cc163 const DeviceAssignment& device_assignment) { in set_device_assignment()
Dexecutable_build_options.h123 const DeviceAssignment& device_assignment() const { in device_assignment() function
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dlibtpu_client.c117 DeviceAssignment device_assignment = {NULL, 0}; in main() local
Dtpu_client.cc522 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
/external/tensorflow/tensorflow/dtensor/mlir/utils/
Dupdate_tpu_metadata.cc124 xla::DeviceAssignmentProto device_assignment; in UpdateTPUCompileMetadata() local
/external/tensorflow/tensorflow/compiler/jit/
Dxla_compile_on_demand_op.cc83 xla::DeviceAssignment device_assignment; in Run() local
/external/tensorflow/tensorflow/core/protobuf/tpu/
Dcompile_metadata.proto90 xla.DeviceAssignmentProto device_assignment = 8; field
/external/tensorflow/tensorflow/python/distribute/
Dtpu_strategy.py682 device_assignment=None): argument
743 device_assignment=None): argument
846 device_assignment=None, argument

123