Searched refs:device_assn (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | collective_ops_test.cc | 312 DeviceAssignment device_assn = MakeDeviceAssn(devices); in XLA_TEST_F() local
316 config.set_static_device_assignment(device_assn); in XLA_TEST_F()
328 /*num_replicas=*/devices.size(), &device_assn, in XLA_TEST_F()
345 DeviceAssignment device_assn; in XLA_TEST_F() member
354 e.device_assn = MakeDeviceAssn(devices); in XLA_TEST_F()
358 config.set_static_device_assignment(e.device_assn); in XLA_TEST_F()
375 .ExecuteReplicated(e.executable.get(), e.opts, &e.device_assn) in XLA_TEST_F()
424 auto device_assn = MakeDeviceAssn(devices); in XLA_TEST_F() local
437 test_runner_.ExecuteReplicated(executable.get(), opts, &device_assn) in XLA_TEST_F()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | replica_id_thunk.cc | 29 TF_ASSIGN_OR_RETURN(auto logical_ids, params.device_assn->LogicalIdsForDevice( in ExecuteOnStream()
|
| D | collective_permute_thunk.cc | 236 RendezvousKey key{params.run_id, params.device_assn->replica_count()}; in ExecuteOnStream()
246 params.device_assn->ReplicaIdForDevice(global_device_id)); in ExecuteOnStream()
|
| D | nccl_collective_thunk.cc | 75 GetParticipatingDevices(global_device_id, *params.device_assn, in ExecuteOnStream()
|
| D | thunk.h | 110 const DeviceAssignment* device_assn; // never null member
|