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 { BasicDeviceList::Create({}), MemoryKind(), /*is_fully_replicated=*/false) {} + MockSharding(tsl::RCReference devices, MemoryKind memory_kind, + bool is_fully_replicated) + : llvm::RTTIExtends(devices, memory_kind, + is_fully_replicated) {} + MOCK_METHOD( (absl::StatusOr< std::vector>>>),