diff --git a/jaxlib/sdy_mpmd.cc b/jaxlib/sdy_mpmd.cc
index 7699a5623946..65190386b6cc 100644
--- a/jaxlib/sdy_mpmd.cc
+++ b/jaxlib/sdy_mpmd.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include <cstdint>
+#include <map>
 #include <optional>
 #include <stdexcept>
 #include <string>