Align Typing with Exported TF Types

[seanpmorgan]

As these become more available we should use the publically exported TF types:
https://github.com/tensorflow/community/blob/master/rfcs/20200211-tf-types.md
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/types/core.py

This is more for standardization than anything so I have no issue waiting until these are mature and we can work through any discrepancies we may have in our types.

[bhack]

See also tensorflow/community#269