Make sure I've backported all moe code from #2371 into the new grug moe from #3046
Make sure I've backported all moe code from #2371 into the new grug moe from #3046