Support NeMo MegatronImagen network

[tfogal]

🚀 Feature

NeMo's "Imagen" network implements text-to-image generation.

Initial examine:

Found 55 distinct operations, of which 46 (83.6%) are supported

Motivation

Pitch

Work items

• [x] #144
• [x] #170
• [ ] #172
• [ ] #173
• [x] #174
• [x] #175
• [x] #176
• [x] #177
• [ ] #178
• [x] Talk to NeMo team about rewriting nan_to_num_ -> nan_to_num (out of place)
• [x] Talk to NeMo team about rewriting add_ -> add (out of place)
• [ ] Test memory use of NeMo model with out-of-place nan_to_num instead of the in-place variant
• [ ] Test memory use of NeMo model with out-of-place add instead of the in-place variant

cc @tfogal

[tfogal]

Quick update: I talked to the NeMo team and Eric had the (reasonable) concern that swapping in-place for out-of-place might increase memory consumption. The onus is on me to give it a try and look at the torch memory use with both in-place and out-of-place options; if the delta is minimal or nonexistent, we could do this.

I'll edit the issue as such.