import numpy as np

import keras
from keras import layers


keras.utils.set_random_seed(7)

device_type = "cpu"
devices = keras.distribution.list_devices(device_type)
if len(devices) < 2:
    raise SystemExit(f"Expected at least 2 CPU devices, found: {devices}")

distribution = keras.distribution.DataParallel(devices=devices)
keras.distribution.set_distribution(distribution)

print(f"Devices: {devices}")
print(f"Distribution: {type(keras.distribution.distribution()).__name__}")

rng = np.random.default_rng(7)
x = rng.normal(size=(128, 4)).astype("float32")
y = (x[:, 0] > 0).astype("float32")

model = keras.Sequential(
    [
        layers.Input(shape=(4,)),
        layers.Dense(1, activation="sigmoid"),
    ]
)
model.compile(
    optimizer=keras.optimizers.SGD(learning_rate=0.2),
    loss="binary_crossentropy",
    metrics=["accuracy"],
)

history = model.fit(x, y, epochs=4, batch_size=16, verbose=2)
print(f"Final accuracy: {history.history['accuracy'][-1]:.4f}")

keras.distribution.set_distribution(None)