Model-parallel Fourier neural operators as learned surrogates for large-scale parametric PDEs

Thomas J. Grady II, Rishi Khan, Mathias Louboutin, Ziyi Yin, Philipp A. Witte, Ranveer Chandra, Russell J. Hewett, Felix J. Herrmann

Abstract: Highlights•Novel approach to parallelism in Fourier neural operators using domain decomposition.•Mathematical description of parallelism in the context of operator learning.•Distributed training of a learned surrogate for subsurface two-phase flow.•Largest forward inference and gradient computation of a Fourier neural operator.