import tensorflow as tf

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    inputs = tf.keras.Input(shape=(4,))
    outputs = tf.keras.layers.Dense(1)(inputs)
    model = tf.keras.Model(inputs, outputs)
    model.compile(optimizer="adam", loss="mse")

features = tf.ones((8, 4))
labels = tf.ones((8, 1))
loss = model.train_on_batch(features, labels)

print("strategy", strategy.__class__.__name__)
print("replicas", strategy.num_replicas_in_sync)
print("model_strategy", model.distribute_strategy is strategy)
print("train_on_batch", loss >= 0)