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

• Support memory_format on to() #144
• Support torch.bernoulli operator. #170
• Implement TensorBase.is_contiguous #172
• TensorBase.repeat_interleave #173
• mse_loss #174
• TensorBase.cuda #175
• Implement nan_to_num #176
• Implement Tensor.type #177
• Implement conv_transpose2d #178
• Talk to NeMo team about rewriting nan_to_num_ -> nan_to_num (out of place)
• 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.