Searched defs:local_device_ids (Results 1 – 6 of 6) sorted by relevance
53 repeated int64 local_device_ids = 4; field
142 const llvm::ArrayRef<int64_t>& local_device_ids, mlir::Value device_id, in DeviceIdToLocalBranchIndex()367 llvm::SmallVector<int64_t> local_device_ids(mesh.local_device_ids().begin(), in ConditionalSave() local640 llvm::SmallVector<int64_t> local_device_ids(mesh.local_device_ids().begin(), in ExpandRestoreV2OpHelper() local
139 absl::Span<const int64_t> local_device_ids() const { in local_device_ids() function
254 const std::vector<std::int64_t>& local_device_ids, in GetMesh()536 std::vector<int64_t> local_device_ids; in GenerateMeshDevicesForTests() local593 std::vector<int64_t> local_device_ids; in FromString() local
90 def _create_device_array(shape, device_type, host_id, local_device_ids=None): argument
375 def local_device_ids(self) -> List[int]: member in Mesh