Custom trianing with overriding the `fit` method single or multiple devices.
### Current Behaviour?

In the official blog post, [HERE](https://www.tensorflow.org/tutorials/distribute/custom_training), it is discussed about the process for single device and multiple device training using `Keras` API. But I've found it quite ambigous to understand the right procedue when you want to do both, [overriding the `fit` method](https://keras.io/guides/customizing_what_happens_in_fit/). It's like combination of custom training loop + using high-level API (`fit`). For example, in custom training loop, it's suggested as follows 

```
train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
test_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)
```

or ephasize to consider during [loss](https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function) calculation. 

or consider the following also while custom training loop?

```
# `run` replicates the provided computation and runs it
# with the distributed input.
@tf.function
def distributed_train_step(dataset_inputs):
  per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
  return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

@tf.function
def distributed_test_step(dataset_inputs):
  return strategy.run(test_step, args=(dataset_inputs,))
```

And it is not clear if we need to do when we combine both (cutom training + fit method).

There are many, for example [this](https://github.com/keras-team/tf-keras/issues/301), that is, if mixed precision is enabled, should we use `LossScaleOptimizer` and `optimizer.get_scaled_loss(loss)`  and `optimizer.get_unscaled_gradients(gradients)` in the `train_step` or `compile` method would do the job? But the [official documentation](https://www.tensorflow.org/api_docs/python/tf/keras/mixed_precision/LossScaleOptimizer) talks about normal fit and custom loop training cases. In case of custom loop, it's suggested to wrap the optimizer and scale the loss and gradient but what about the combination of fit and custom loop (overriding train_step)? Does it sill need to wrap the optimizer and scale the loss and gradient or it will be handled by the API?


### Relevant log output