FLAX&JAX Implementation of the paper: "Neural Discrete Representation Learning"