JAXify

This may require changes to the selection function interpolation or representation, so that it can be differentiable.  Otherwise the likelihood itself and the DF models should be easy to implement in JAX.