*   Metrics update and collection logic in default `Model.train_step()` is
        now customizable via overriding `Model.compute_metrics()`.
    *   Losses computation logic in default `Model.train_step()` is now
        customizable via overriding `Model.compute_loss()`.
    *   `jit_compile` added to `Model.compile()` on an opt-in basis to compile