Modify core encode decode functions to pass a dict to the prediction function

[JossWhittle]

Currently the lowres inputs are passed directly to the prediction_fn as the only input.

• Modify to accept a dict that has at least one key for the lowres input.

• Provide boolean flag to also pass a positional encoding tensor along with the lowres which the model can use if needed.

• Chunked encode decode will need to generate the correct chunks of the positional encoding for the current chunk.

• Model can choose how to use positional encodings.

  • Image case would receive (B, H, W, 2) tensor containing the Y and X coordinates of each pixel in the trailing axis.
  • Volume case would receive (B, D, H, W, 3) tensor containing the Z, Y, and X coordinates of each voxel in the trailing axis.