From discussion with @qiencai @radka-j @ContiPaolo, a useful extension would be to implement support for masks.

This might be in a few different contexts:

• Fixed mask (both at training and inference time, e.g. a land mass)
  • In this case it could be passed upon initialization of the model and used in the loss and forward calculations
• Unknown mask at inference time (e.g. acoustic scattering - maze in The Well)
  • In this case the specific mask would need to be passed at inference time
  • We might look to extend the Batch API to have a field called mask and this could be populated/derived from a given constant field already present or some additional data

Two example papers on this:

• Palette: Image-to-Image Diffusion Models (loss and sampling process both exclude masked pixels)
• RePaint: Inpainting using Denoising Diffusion Probabilistic Models (no mask during training/loss but at sampling the mask is also noised but refreshed every sample step. There's also a Azula sampler for this.

Other thoughts:

• If we can identify API that can be flexible and minimise repeated implementation would be great.
• API to support mask concatenation as a new channel might be worthwhile.
• An outstanding question is how to handle masking that works in the latent space too.