Utility to run tests inside tf.function and eager.

Relies on being able to run the assert* test methods inside a
py_func to run them inside the graph, so there's no need for
self.evaluate or similar methods which create a graph/eager
hybrid programming model.

PiperOrigin-RevId: 224575790