// Explicit Tensor precision and placement; no silent device selection.
import aner.nn
import aner.dev

fn exercise(device: String) -> Tensor {
    nn.require_device(device)
    let input = nn.random(32, 64, seed: 2026, scale: 0.25)
        .astype(dtype: "Float32").to(device: device)
    var weights = nn.parameter(nn.random(64, 16, seed: 42, scale: 0.1)
        .astype(dtype: "Float32").to(device: device))
    let target = nn.zeros(32, 16).astype(dtype: "Float32").to(device: device)
    let first = nn.mse(nn.matmul(input, weights), target)
    print(device)
    print(weights.dtype())
    print(weights.backend())
    print(nn.item(first))
    var step = 0
    while step < 8 {
        let loss = nn.mse(nn.matmul(input, weights), target)
        nn.backward(loss)
        weights = nn.sgd(weights, rate: 0.1)
        step = step + 1
    }
    nn.synchronize(device)
    let final = nn.mse(nn.matmul(input, weights), target)
    print(nn.item(final))
    dev.capture("resident trained weights", weights)
    let portable = nn.detach(weights).to(device: "cpu")
    print(portable.device())
    print(portable.dtype())
    return nn.detach(weights)
}

print(nn.devices())
let cpu_weights = exercise("cpu")
var mps_weights = nn.zeros(1, 1).astype(dtype: "Float32")
var cuda_weights = nn.zeros(1, 1).astype(dtype: "Float32")
if nn.device_available("mps") { mps_weights = exercise("mps") }
else { print("MPS unavailable: this build or device does not support it") }
if nn.device_available("cuda") { cuda_weights = exercise("cuda") }
else { print("CUDA unavailable: this build or device does not support it") }