https://arxiv.org/abs/1809.00738 is really nice for general context. Plus can implement some of these generalizations from what is already here.