From bc1aad8a0de112d28b16b14b8dde1645b6ede52b Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Tue, 17 Sep 2024 14:11:14 -0700
Subject: [PATCH] Add constructor to MockSharding that also takes device list
 to allow tests that need working call to Sharding::devices().

PiperOrigin-RevId: 675704343
---
 xla/python/ifrt/mock.h | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/xla/python/ifrt/mock.h b/xla/python/ifrt/mock.h
index 77209214543d1..125c95e58a3fb 100644
--- a/xla/python/ifrt/mock.h
+++ b/xla/python/ifrt/mock.h
@@ -324,6 +324,11 @@ class MockSharding : public llvm::RTTIExtends<MockSharding, Sharding> {
             BasicDeviceList::Create({}), MemoryKind(),
             /*is_fully_replicated=*/false) {}
 
+  MockSharding(tsl::RCReference<DeviceList> devices, MemoryKind memory_kind,
+               bool is_fully_replicated)
+      : llvm::RTTIExtends<MockSharding, Sharding>(devices, memory_kind,
+                                                  is_fully_replicated) {}
+
   MOCK_METHOD(
       (absl::StatusOr<
           std::vector<std::pair<Shape, std::shared_ptr<const Sharding>>>>),