// CPU and available accelerator classifiers use the same data and architecture.
// CPU uses Float64; GPU uses Float32. Compare metrics within numerical tolerance.
// For matched Float32 Tensor throughput, use the native aner_gpu_benchmark tool.
import aner.dataset
import aner.nn
import aner.ml
import aner.dev

fn train(device: String, split: DatasetSplit, network: NeuralNetwork) {
    let plan = ml.classifier(network, standardize: true,
        loss: "cross_entropy", optimizer: "adam").on(device: device)
    let fitted = plan.fit(split.train(), epochs: 20, rate: 0.02, progress: true)
    print(device)
    print(fitted.weights(0).device())
    print(fitted.weights(0).dtype())
    print(fitted.weights(0).backend())
    print(fitted.training_seconds())
    print(fitted.initial_loss())
    print(fitted.training_loss())
    print(fitted.evaluate(split.test()).accuracy())
    dev.capture("trained classifier", fitted)
    dev.capture("resident classifier weights", fitted.weights(0))
}

print(nn.devices())
let split = dataset.iris().split(test: 0.2, seed: 2026)
let network = nn.sequential(seed: 42)
    .dense(inputs: 4, outputs: 16, initialization: "he").relu()
    .dense(inputs: 16, outputs: 3, initialization: "xavier")
train("cpu", split, network)
if nn.device_available("mps") { train("mps", split, network) }
if nn.device_available("cuda") { train("cuda", split, network) }