Add tf.function support

[schneimo]

To add tf.function support some parts of the main code must be heavily refactored since else the project won't work anymore or it will be slower than in eager mode.