// Bounded Tensor stress. Use aner_gpu_benchmark for synchronized measurements.
// Increase repetitions deliberately; dimensions stay below Tensor/work limits.
import aner.nn
import aner.dev

fn stress(device: String) {
    nn.require_device(device)
    let left = nn.random(256, 256, seed: 2026, scale: 0.1)
        .astype(dtype: "Float32").to(device: device)
    let right = nn.random(256, 256, seed: 2027, scale: 0.1)
        .astype(dtype: "Float32").to(device: device)
    var output = nn.matmul(left, right)
    var iteration = 0
    while iteration < 10 {
        output = nn.matmul(left, right)
        iteration = iteration + 1
    }
    nn.synchronize(device)
    print(device)
    print(output.dtype())
    print(output.backend())
    print(nn.item(nn.mean(output)))
    dev.capture("resident matrix product", output)
}

print(nn.devices())
stress("cpu")
if nn.device_available("mps") { stress("mps") }
if nn.device_available("cuda") { stress("cuda") }