How to extract a relay function when doing partitioning with relay pattern?

Hi @comaniac , @masahi and @jroesch,

Any suggestions for the above problem? How can I extract the partitioned relay function from relay IRModule?