Rewrite TensorFlow train_step and test_step
Draft PR for a full rewrite of the TF train/test steps. I swear this will fix like 50% of our TF issues in one PR.

Current status:
- Correctly handles output mapping across most model classes for losses + metrics
- Keras metrics are back, even with the dummy loss. (!!!!)
- Keras metrics work correctly even for multi-output models (like QA)
- In most cases, users can pass tensors in either the input dict or the labels and the model will handle them correctly.
- No more errors when calling `fit()` when the model has nested output structure (e.g. the model outputting a `past` tuple)

What's left to do:

- [X] Models with multiple unusual outputs that do not match label names may still have issues with metrics. This is relatively uncommon. We support adding a property to those classes to tell Keras what to do with the labels, but we haven't added it to any models yet. (None are failing in tests, so hopefully we won't need to worry too much about this!)
- [x] Testing testing testing! I want to rerun all notebooks/examples and make sure the user experience is good.
- [X] CI testing - We need to make sure we don't regress on any of this
- [ ] Discoverability: After this is merged we should update notebooks/examples to show off the cool new features, and document our TF workflow/philosophy somewhere that new users will find.

_The documentation is not available anymore as the PR was closed or merged._<|||||>(Requesting reviews now that @gante is back)