Add constructor to MockSharding that also takes device list to allow …

…tests that need working call to Sharding::devices(). PiperOrigin-RevId: 674504128
openxla · Sep 17, 2024 · 2ba7753 · 2ba7753
1 parent 8ace4ee
commit 2ba7753
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/xla/python/ifrt/mock.h b/xla/python/ifrt/mock.h
@@ -324,6 +324,12 @@ class MockSharding : public llvm::RTTIExtends<MockSharding, Sharding> {
             BasicDeviceList::Create({}), MemoryKind(),
             /*is_fully_replicated=*/false) {}
 
+  explicit MockSharding(tsl::RCReference<DeviceList> devices,
+                        MemoryKind memory_kind = MemoryKind(),
+                        bool is_fully_replicated = false)
+      : llvm::RTTIExtends<MockSharding, Sharding>(devices, memory_kind,
+                                                  is_fully_replicated) {}
+
   MOCK_METHOD(
       (absl::StatusOr<
           std::vector<std::pair<Shape, std::shared_ptr<const Sharding>>>>),