// Select the execution profile with CLI flags or nn.configure in a notebook.
// Default: CPU Float64. Use --device auto for an available accelerator.
import aner.nn
import aner.dev

print(nn.execution())
let source: Tensor = [[1.00000001]]
let placed = source.place()
print(placed.dtype())
print(placed.device())
print(placed.backend())
print(nn.value(placed, row: 0, column: 0))

let raw_input: Tensor = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
let raw_target: Tensor = [[1.0, 0.0], [0.0, 1.0]]
let input = raw_input.place()
let target = raw_target.place()
var weights = nn.parameter(nn.random(rows: 3, cols: 2, seed: 42, scale: 0.1).place())
var first_loss = 0.0
var epoch = 0
while epoch < 8 {
    let prediction = nn.matmul(input, weights)
    let loss = nn.mse(prediction, target)
    if epoch == 0 {
        first_loss = nn.item(loss)
    }
    nn.backward(loss)
    dev.capture(label: "precision training", value: loss)
    weights = nn.sgd(weights, rate: 0.1)
    epoch = epoch + 1
}
let final_loss = nn.mse(nn.matmul(input, weights), target)
nn.synchronize(weights.device())
print(first_loss)
print(nn.item(final_loss))
print(weights.dtype())
print(weights.device())
