Port to tensorflow 2

We need to port this package to tensorflow 2 and remove unused parts.