Hello there @purpledog! Just a friendly FYI: We at IBM's CODAIT center have been working on a new version of your code in TensorFlow's contrib.graph_editor module. See the parent project of this issue for more information. Please feel free to contact us if you have any questions or would like to be involved!