Closes #14. Functions which currently return an ordered tuple of maps (lrmap, udmap, cmap, ...) now return keyed dictionaries { 'lrmap': lrmap, 'udmap': udmap, 'cmap': cmap, ... } so that order/usage is explicitly enforced.

List comprehensions over the tuples now use jax.tree_map and jax.tree_multimap to ensure key safety.

@GMW99, this will break the current implementation of the Metrics Callback class which iterates over a zip of the hardcoded map names and the maps tuple. This iteration can be replaced by iterating over maps.items() since it is now a dict already.

Functions that currently mark static_argnums=(0, 1, 2) should be updated to use the safer static_argnames=('tom', 'dick', 'harry') that is now available.

• Splits docker image into JAX base image and Kompressor dependency and install image
• JAX image installs JAX from source to ensure correct CUDA / CUDNN versions
• Adjust setup.py to install dependencies from requirement.txt
• Refactors a how submodules are imported (within the kom.image submodule. Need to check volumes matches)
• Add kom.image.data submodule for dealing with tensorflow data pipelines
• Fixed pooling in the total variation losses (used as metrics in the example notebooks)
• Move all the encoding/decoding functions for the maps into a kom.mapping submodule
• Add within-k and run-length metrics to kom.image.metrics for example notebooks
• Added example notebooks for interacting with the maps and training a basic Haiku compression model

Will need a bit of thinking to get right. We probably need to consider similar tricks that we used for applying Radix Sort on float32 data to make the compression numerically stable and portable between machines.

Some of our data is uint32 volumes.

Will need to trace through the full compression implementation and make sure intermediate value dtypes are large enough to avoid uint32 overflow when needed.

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.

Decompress sliced chunk of image or volume without needing to decompress the entire data element.

• May require applying secondary compression in blocks to avoid needing to decompress the full level maps, only to apply the predictor to the target slice.

• Instead unpack just the blocks needed for the slice then trim.

• A kompressor (or stack of) trained to secondary compress the maps from the primary kompressor (or stack of) would be able to naturally handle slice chunked decoding.

  • Could such a secondary compressor be shared between levels? Between multiple kompressors in the primary stack?

• Experiment with implementing the 1D case for compressing signals.
• Video as sequence of 2D frames using the 3D volume code directly.
• Look at compressing within timestep using information from neighbouring timesteps without actually compressing (dropping frames) the temporal axis.