Searched refs:device_assignment_ (Results 1 – 8 of 8) sorted by relevance
105 bool has_device_assignment() const { return device_assignment_.has_value(); } in has_device_assignment()109 CHECK(device_assignment_.has_value()); in device_assignment()110 return device_assignment_.value(); in device_assignment()152 absl::optional<DeviceAssignment> device_assignment_; variable
94 device_assignment_ = device_assignment; in set_device_assignment()
97 device_assignment_ = device_assignment; in set_device_assignment()102 return device_assignment_; in device_assignment()
161 const DeviceAssignment* device_assignment_ = nullptr; variable
314 int num_replicas() const { return device_assignment_.replica_count(); } in num_replicas()315 int num_partitions() const { return device_assignment_.computation_count(); } in num_partitions()323 return device_assignment_; in device_assignment()367 const DeviceAssignment device_assignment_; variable
514 device_assignment_(std::move(device_assignment)), in PyTpuExecutable()517 VLOG(1) << "DeviceAssignment. " << device_assignment_.ToString(); in PyTpuExecutable()518 const int num_replicas = device_assignment_.replica_count(); in PyTpuExecutable()519 const int num_partitions = device_assignment_.computation_count(); in PyTpuExecutable()523 int device_id = device_assignment_(replica, partition); in PyTpuExecutable()548 const int device_id = device_assignment_(replica, partition); in ExecuteHelper()582 CHECK(device_assignment_.Serialize(&device_assignment).ok()); in ExecuteHelper()
1574 device_assignment_(std::move(device_assignment)), in PjRtStreamExecutorExecutable()1585 if (device_assignment_ == nullptr) { in PjRtStreamExecutorExecutable()1593 << device_assignment_->ToString(); in PjRtStreamExecutorExecutable()1594 CHECK_GE(addressable_devices_.size(), 1) << device_assignment_->ToString(); in PjRtStreamExecutorExecutable()1597 num_partitions = device_assignment_->computation_count(); in PjRtStreamExecutorExecutable()1863 CHECK(device_assignment_ != nullptr); in ExecuteHelper()1864 const int device_id = (*device_assignment_)(replica, partition); in ExecuteHelper()1866 device_assignment = device_assignment_; in ExecuteHelper()1868 CHECK(device_assignment_ == nullptr); in ExecuteHelper()1945 if (device_assignment_ == nullptr) { in Execute()[all …]
634 return *device_assignment_; in device_assignment()715 std::shared_ptr<DeviceAssignment> device_assignment_; variable