Imported from https://github.com/Slowpuncher24/pde-net-in-tf