Library being used for various functions in the notebook.

Thanks for the great tutorial and paper, I see a few pieces are missing from the code, from which library are we importing the class Ptree() from, I don't see any imports that refer to this. Is this some internal library(or causal inference engine) being used by deepmind or is it available for public use? I am doing a project on similar grounds and such a library would be of great help(although I know I can do all sorts of experimentation within the colab notebook provided, but was just curious about it).