---
title: "Keras-like API in ggmlR"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{02. Keras-like API in ggmlR}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include=FALSE}
# Vignette code is executed locally (NOT_CRAN=true) but not on CRAN, where
# the CPU fallback would multi-thread and trip the "CPU time > elapsed" NOTE.
knitr::opts_chunk$set(eval = identical(Sys.getenv("NOT_CRAN"), "true"))
```

ggmlR provides a high-level Keras-like interface for building and training neural
networks on top of the ggml C tensor library. This vignette walks through the
main APIs using the built-in `iris` and `mtcars` datasets.

## Installation

```{r}
# install.packages("ggmlR")   # once on CRAN
library(ggmlR)
```

---

## 1. Sequential API

The sequential API mirrors `keras.Sequential`: stack layers with `|>`, compile,
fit, evaluate, predict.

### 1.1 Prepare data — iris (3-class classification)

```{r}
data(iris)
set.seed(42)

x <- as.matrix(iris[, 1:4])
x <- scale(x)                        # standardise

# one-hot encode species
y <- model.matrix(~ Species - 1, iris)  # [150 x 3]

idx   <- sample(nrow(x))
n_tr  <- 120L
x_tr  <- x[idx[1:n_tr],  ]
y_tr  <- y[idx[1:n_tr],  ]
x_val <- x[idx[(n_tr+1):150], ]
y_val <- y[idx[(n_tr+1):150], ]
```

### 1.2 Build and compile

```{r}
model <- ggml_model_sequential() |>
  ggml_layer_dense(32L, activation = "relu", input_shape = 4L) |>
  ggml_layer_batch_norm() |>
  ggml_layer_dropout(0.3, stochastic = TRUE) |>
  ggml_layer_dense(16L, activation = "relu") |>
  ggml_layer_dense(3L,  activation = "softmax") |>
  ggml_compile(
    optimizer = "adam",
    loss      = "categorical_crossentropy",
    metrics   = c("accuracy")
  )

print(model)
```

Available losses:

| `loss =` | Notes |
|---|---|
| `"categorical_crossentropy"` | Multi-class; expects a `softmax` output |
| `"mse"` | Regression, squared error |
| `"mae"` | Regression, absolute error — robust to outliers |
| `"huber"` | Quadratic near zero, linear beyond `delta = 1` |
| `"binary_crossentropy"` | Independent per-output labels; expects a `sigmoid` output |

`"mean_squared_error"`, `"mean_absolute_error"`, `"huber_loss"` and
`"binary_cross_entropy"` are accepted as aliases.

MAE and Huber matter when the targets contain outliers: MSE squares the error,
so one wild value dominates the gradient, while MAE weights every point equally
and Huber is quadratic only near zero. Binary cross-entropy treats each output
as its own Bernoulli, which is what makes it the loss for multi-label targets
rather than for one-of-N classification.

### 1.3 Train

```{r}
model <- ggml_fit(
  model,
  x_tr, y_tr,
  epochs           = 100L,
  batch_size       = 32L,
  validation_split = 0.0,   # we supply our own val set below
  verbose          = 0L
)
```

### 1.4 Evaluate and predict

```{r}
score <- ggml_evaluate(model, x_val, y_val, batch_size = 16L)
cat(sprintf("Val loss: %.4f   Val accuracy: %.4f\n", score$loss, score$accuracy))

probs   <- ggml_predict(model, x_val, batch_size = 16L)
classes <- apply(probs, 1, which.max)
true    <- apply(y_val, 1, which.max)
cat("Confusion matrix:\n")
print(table(true = true, predicted = classes))
```

### 1.5 Save and reload

```{r}
path <- tempfile(fileext = ".rds")
ggml_save_model(model, path)
cat(sprintf("Saved: %s (%.1f KB)\n", path, file.size(path) / 1024))

model2 <- ggml_load_model(path)
score2 <- ggml_evaluate(model2, x_val, y_val, batch_size = 16L)
cat(sprintf("Reloaded model accuracy: %.4f\n", score2$accuracy))
```

---

## 2. Functional API

The functional API lets you build DAGs with multiple inputs/branches/outputs.
Use `ggml_input()` to declare inputs, pipe through layers, then call `ggml_model()`.

### 2.1 Single-input DAG — residual-style network on iris

```{r}
inp  <- ggml_input(shape = 4L)

h    <- inp |> ggml_layer_dense(32L, activation = "relu") |> ggml_layer_batch_norm()
h    <- h   |> ggml_layer_dense(16L, activation = "relu")
out  <- h   |> ggml_layer_dense(3L,  activation = "softmax")

model_fn <- ggml_model(inputs = inp, outputs = out) |>
  ggml_compile(optimizer = "adam", loss = "categorical_crossentropy",
               metrics = c("accuracy"))

model_fn <- ggml_fit(model_fn, x_tr, y_tr,
                     epochs = 100L, batch_size = 32L, verbose = 0L)

score_fn <- ggml_evaluate(model_fn, x_val, y_val, batch_size = 16L)
cat(sprintf("Functional model accuracy: %.4f\n", score_fn$accuracy))
```

### 2.2 Multi-input model — mtcars regression

Predict `mpg` from two input groups: engine features and transmission features.

```{r}
data(mtcars)
set.seed(7)

# Input group 1: engine (disp, hp, wt)
# Input group 2: transmission / gearbox (cyl, gear, carb, am)
engine <- as.matrix(scale(mtcars[, c("disp","hp","wt")]))
trans  <- as.matrix(scale(mtcars[, c("cyl","gear","carb","am")]))
y_mpg  <- matrix(scale(mtcars$mpg), ncol = 1L)          # [32 x 1]

# small dataset — use all for training, evaluate on same data for demo
x1 <- engine;  x2 <- trans

inp1 <- ggml_input(shape = 3L, name = "engine")
inp2 <- ggml_input(shape = 4L, name = "transmission")

branch1 <- inp1 |> ggml_layer_dense(16L, activation = "relu")
branch2 <- inp2 |> ggml_layer_dense(16L, activation = "relu")

merged  <- ggml_layer_add(list(branch1, branch2))        # element-wise add
out_reg <- merged |>
  ggml_layer_dense(8L, activation = "relu") |>
  ggml_layer_dense(1L)

model_reg <- ggml_model(inputs = list(inp1, inp2), outputs = out_reg) |>
  ggml_compile(optimizer = "adam", loss = "mse")

model_reg <- ggml_fit(model_reg,
                      x = list(x1, x2), y = y_mpg,
                      epochs = 200L, batch_size = 16L, verbose = 0L)

preds <- ggml_predict(model_reg, x = list(x1, x2), batch_size = 16L)
cat(sprintf("Pearson r (scaled mpg): %.4f\n", cor(preds, y_mpg)))
```

`ggml_layer_add()` requires the branches to have the same width.
`ggml_layer_concatenate()` does not — it stacks them along an axis, keeping each
branch's features distinct instead of summing them. Both are trainable:
gradients flow back into every branch.

```{r}
br1 <- inp1 |> ggml_layer_dense(8L,  activation = "relu")
br2 <- inp2 |> ggml_layer_dense(12L, activation = "relu")   # widths differ

out_cat <- ggml_layer_concatenate(list(br1, br2), axis = 0L) |>
  ggml_layer_dense(1L)

model_cat <- ggml_model(inputs = list(inp1, inp2), outputs = out_cat) |>
  ggml_compile(optimizer = "adam", loss = "mse")
model_cat <- ggml_fit(model_cat, x = list(x1, x2), y = y_mpg,
                      epochs = 200L, batch_size = 16L, verbose = 0L)

p_cat <- ggml_predict(model_cat, x = list(x1, x2), batch_size = 16L)
cat(sprintf("Concatenate variant, Pearson r: %.4f\n", cor(p_cat, y_mpg)))
```

### 2.3 Multi-output model — one trunk, two heads

A model with several outputs trains *all* of them. Each head gets its own loss
and weight, and the optimizer minimizes `sum(loss_weights[i] * loss_i)`. Losses
may be mixed: below a classification head and a regression head share one trunk.

Pass `y` as a **list**, one matrix per head, matched by output name.

```{r}
set.seed(11)

# Two targets from the same cars: efficient or not (class), and mpg (value).
y_cls <- as.integer(mtcars$mpg > median(mtcars$mpg))
y_cls <- cbind(1 - y_cls, y_cls) * 1.0         # one-hot [32 x 2]
x_all <- as.matrix(scale(mtcars[, c("disp","hp","wt","cyl")]))

inp_mo <- ggml_input(shape = 4L, name = "features")
trunk  <- inp_mo |> ggml_layer_dense(16L, activation = "relu", name = "trunk")
head_c <- trunk  |> ggml_layer_dense(2L, activation = "softmax", name = "class")
head_v <- trunk  |> ggml_layer_dense(1L, name = "value")

model_mo <- ggml_model(inputs = inp_mo, outputs = list(head_c, head_v)) |>
  ggml_compile(optimizer = "adam",
               loss         = list(class = "categorical_crossentropy",
                                   value = "mse"),
               loss_weights = c(class = 1.0, value = 0.5))

model_mo <- ggml_fit(model_mo, x_all,
                     y = list(class = y_cls, value = y_mpg),
                     epochs = 150L, batch_size = 16L, verbose = 0L)

# Per-head losses, so a head that stops learning is visible on its own
# rather than hidden inside the total.
h <- model_mo$history
cat(sprintf("class loss: %.4f -> %.4f\n",
            h$train_class_loss[1], tail(h$train_class_loss, 1)))
cat(sprintf("value loss: %.4f -> %.4f\n",
            h$train_value_loss[1], tail(h$train_value_loss, 1)))
```

`ggml_predict()` returns a list, one matrix per head, and `ggml_evaluate()`
reports a `<output>_loss` for each.

### 2.4 Attention and a transformer block

`ggml_layer_attention()` is multi-head scaled dot-product attention. It takes
and returns a sequence node of shape `c(seq_len, d_model)`, so blocks stack and
a residual merge needs no reshaping.

`time_distributed = TRUE` applies one dense kernel at every position, sharing
weights across the sequence — the position-wise feed-forward sublayer that
completes a transformer block. (Without it a dense layer would flatten the
sequence into a single vector.)

```{r}
set.seed(13)
n_seq <- 128L; S <- 6L; D <- 8L

# A task that needs mixing across positions: is the sequence sum positive?
xs   <- array(runif(n_seq * S * D, -1, 1), dim = c(n_seq, S, D))
lab  <- as.integer(apply(xs, 1, mean) > 0)
y_sq <- cbind(1 - lab, lab) * 1.0

inp_t <- ggml_input(shape = c(S, D), name = "seq")

# Attention sublayer + residual
att <- inp_t |> ggml_layer_attention(d_model = D, n_heads = 2L, name = "mha")
h1  <- ggml_layer_add(list(inp_t, att))

# Position-wise feed-forward sublayer + residual
ff  <- h1 |> ggml_layer_dense(D * 2L, activation = "relu",
                              time_distributed = TRUE, name = "ff1")
ff  <- ff |> ggml_layer_dense(D, time_distributed = TRUE, name = "ff2")
h2  <- ggml_layer_add(list(h1, ff))

out_t <- h2 |> ggml_layer_flatten() |>
  ggml_layer_dense(2L, activation = "softmax", name = "cls")

model_t <- ggml_model(inputs = inp_t, outputs = out_t) |>
  ggml_compile(optimizer = "adam", loss = "categorical_crossentropy")

model_t <- ggml_fit(model_t, xs, y_sq, epochs = 30L, batch_size = 32L,
                    verbose = 0L)
cat(sprintf("transformer block accuracy: %.4f\n",
            tail(model_t$history$train_accuracy, 1)))
```

`causal = TRUE` masks keys after the query, making the layer autoregressive
(GPT-style). Cross-attention takes queries from one node and keys/values from
another, which may have a different length:

```{r eval=FALSE}
dec <- inp_t |> ggml_layer_attention(D, n_heads = 2L, causal = TRUE)

ctx_seq <- ggml_input(shape = c(10L, D))
x_cross <- ggml_apply(list(dec, ctx_seq), ggml_attention(D, n_heads = 2L))
```

`ggml_attention()` returns a reusable layer object: applying it to several nodes
shares one set of projections, which is how an encoder block is reused across a
stack. All heads are computed in one batched pass, so the graph does not grow
with `n_heads`.

---

## 3. Callbacks

Callbacks plug into `ggml_fit()` via the `callbacks` argument.

### 3.1 Early stopping

```{r}
cb_stop <- ggml_callback_early_stopping(
  monitor   = "val_loss",
  patience  = 15L,
  min_delta = 1e-4
)

model_cb <- ggml_model_sequential() |>
  ggml_layer_dense(32L, activation = "relu", input_shape = 4L) |>
  ggml_layer_dense(3L,  activation = "softmax") |>
  ggml_compile(optimizer = "adam", loss = "categorical_crossentropy")

model_cb <- ggml_fit(model_cb, x_tr, y_tr,
                     epochs           = 300L,
                     batch_size       = 32L,
                     validation_split = 0.1,
                     callbacks        = list(cb_stop),
                     verbose          = 0L)
```

### 3.2 Learning-rate schedulers

```{r}
# Cosine annealing
cb_cosine <- ggml_schedule_cosine_decay(T_max = 100L, eta_min = 1e-5)

# Step decay: halve LR every 30 epochs
cb_step <- ggml_schedule_step_decay(step_size = 30L, gamma = 0.5)

# Reduce on plateau
cb_plateau <- ggml_schedule_reduce_on_plateau(
  monitor  = "val_loss",
  factor   = 0.5,
  patience = 10L,
  min_lr   = 1e-6
)

model_lr <- ggml_model_sequential() |>
  ggml_layer_dense(32L, activation = "relu", input_shape = 4L) |>
  ggml_layer_dense(3L,  activation = "softmax") |>
  ggml_compile(optimizer = "adam", loss = "categorical_crossentropy")

model_lr <- ggml_fit(model_lr, x_tr, y_tr,
                     epochs           = 150L,
                     batch_size       = 32L,
                     validation_split = 0.1,
                     callbacks        = list(cb_cosine),
                     verbose          = 0L)
```

---

## 4. Autograd Engine (`ag_*`)

The autograd engine provides PyTorch-style dynamic computation graphs.
Tensors are **column-major**: shape is `[features, batch]`.

### 4.1 Prepare data (iris, col-major)

```{r}
# transpose: rows = features, cols = samples
x_tr_ag  <- t(x_tr)    # [4, 120]
y_tr_ag  <- t(y_tr)    # [3, 120]
x_val_ag <- t(x_val)   # [4, 30]
y_val_ag <- t(y_val)
```

### 4.2 Build a model with `ag_sequential`

```{r}
ag_mod <- ag_sequential(
  ag_linear(4L,  32L, activation = "relu"),
  ag_batch_norm(32L),
  ag_dropout(0.3),
  ag_linear(32L, 16L, activation = "relu"),
  ag_linear(16L,  3L)
)

params <- ag_mod$parameters()
opt    <- optimizer_adam(params, lr = 1e-3)
```

### 4.3 Training loop

```{r}
BS <- 32L
n  <- ncol(x_tr_ag)

ag_train(ag_mod)
set.seed(42)

for (ep in seq_len(150L)) {
  perm <- sample(n)
  for (b in seq_len(ceiling(n / BS))) {
    idx <- perm[seq((b-1L)*BS + 1L, min(b*BS, n))]
    xb  <- ag_tensor(x_tr_ag[, idx, drop = FALSE])
    yb  <- y_tr_ag[, idx, drop = FALSE]

    with_grad_tape({
      loss <- ag_softmax_cross_entropy_loss(ag_mod$forward(xb), yb)
    })
    grads <- backward(loss)
    opt$step(grads)
    opt$zero_grad()
  }

  if (ep %% 50L == 0L)
    cat(sprintf("epoch %d  loss %.4f\n", ep, loss$data[1]))
}
```

### 4.4 Inference and accuracy

```{r}
ag_eval(ag_mod)

# forward in chunks, apply softmax manually
ag_predict_cm <- function(model, x_cm, chunk = 64L) {
  n   <- ncol(x_cm)
  out <- matrix(0.0, nrow(model$forward(ag_tensor(x_cm[,1,drop=FALSE]))$data), n)
  for (s in seq(1L, n, by = chunk)) {
    e  <- min(s + chunk - 1L, n)
    lg <- model$forward(ag_tensor(x_cm[, s:e, drop = FALSE]))$data
    ev <- exp(lg - apply(lg, 2, max))
    out[, s:e] <- ev / colSums(ev)
  }
  out
}

probs_ag <- ag_predict_cm(ag_mod, x_val_ag)          # [3, 30]
pred_ag  <- apply(probs_ag, 2, which.max)
true_ag  <- apply(y_val_ag, 1, which.max)            # col-major: rows = classes
acc_ag   <- mean(pred_ag == true_ag)
cat(sprintf("Autograd val accuracy: %.4f\n", acc_ag))
```

### 4.5 LR scheduler + gradient clipping

```{r}
ag_mod2 <- ag_sequential(
  ag_linear(4L,  64L, activation = "relu"),
  ag_batch_norm(64L),
  ag_dropout(0.3),
  ag_linear(64L, 32L, activation = "relu"),
  ag_linear(32L,  3L)
)
params2 <- ag_mod2$parameters()
opt2    <- optimizer_adam(params2, lr = 1e-3)
sch2    <- lr_scheduler_cosine(opt2, T_max = 150L, lr_min = 1e-5)

ag_train(ag_mod2)
set.seed(42)

for (ep in seq_len(150L)) {
  perm <- sample(n)
  for (b in seq_len(ceiling(n / BS))) {
    idx <- perm[seq((b-1L)*BS + 1L, min(b*BS, n))]
    xb  <- ag_tensor(x_tr_ag[, idx, drop = FALSE])
    yb  <- y_tr_ag[, idx, drop = FALSE]

    with_grad_tape({
      loss2 <- ag_softmax_cross_entropy_loss(ag_mod2$forward(xb), yb)
    })
    grads2 <- backward(loss2)
    clip_grad_norm(params2, grads2, max_norm = 5.0)
    opt2$step(grads2)
    opt2$zero_grad()
  }
  sch2$step()
}

ag_eval(ag_mod2)
probs2 <- ag_predict_cm(ag_mod2, x_val_ag)
acc2   <- mean(apply(probs2, 2, which.max) == true_ag)
cat(sprintf("ag + cosine + clip  val accuracy: %.4f\n", acc2))
```

### 4.6 Manual model from raw `ag_param`

For full control, build a network directly from parameters:

```{r}
make_net <- function(n_in, n_hidden, n_out) {
  W1 <- ag_param(matrix(rnorm(n_hidden * n_in) * sqrt(2/n_in),  n_hidden, n_in))
  b1 <- ag_param(matrix(0.0, n_hidden, 1L))
  W2 <- ag_param(matrix(rnorm(n_out * n_hidden) * sqrt(2/n_hidden), n_out, n_hidden))
  b2 <- ag_param(matrix(0.0, n_out, 1L))

  list(
    forward    = function(x)
      ag_add(ag_matmul(W2, ag_relu(ag_add(ag_matmul(W1, x), b1))), b2),
    parameters = function() list(W1=W1, b1=b1, W2=W2, b2=b2)
  )
}

set.seed(1)
net    <- make_net(4L, 32L, 3L)
opt_r  <- optimizer_adam(net$parameters(), lr = 1e-3)

for (ep in seq_len(200L)) {
  perm <- sample(n)
  for (b in seq_len(ceiling(n / BS))) {
    idx <- perm[seq((b-1L)*BS+1L, min(b*BS, n))]
    xb  <- ag_tensor(x_tr_ag[, idx, drop = FALSE])
    yb  <- y_tr_ag[, idx, drop = FALSE]
    with_grad_tape({ loss_r <- ag_softmax_cross_entropy_loss(net$forward(xb), yb) })
    gr <- backward(loss_r)
    opt_r$step(gr);  opt_r$zero_grad()
  }
}

probs_r <- ag_predict_cm(net, x_val_ag)
acc_r   <- mean(apply(probs_r, 2, which.max) == true_ag)
cat(sprintf("Raw ag_param val accuracy: %.4f\n", acc_r))
```

---

## 5. GPU / device selection

```{r}
# Use Vulkan GPU if available, fall back to CPU
device <- tryCatch({
  ag_device("gpu")
  "gpu"
}, error = function(e) "cpu")

cat("Running on:", device, "\n")

# Mixed precision (f16 on GPU, f32 on CPU)
ag_dtype(if (device == "gpu") "f16" else "f32")
```

After calling `ag_device()` and `ag_dtype()`, all subsequent `ag_param` and
`ag_tensor` calls use the selected device and dtype.  The sequential API
(`ggml_fit`) picks up the backend automatically.

---

## API summary

| Task | Sequential / Functional | Autograd |
|------|------------------------|----------|
| Build model | `ggml_model_sequential()` + `ggml_layer_*` | `ag_sequential()` or raw `ag_param` |
| Functional DAG | `ggml_input()` → `ggml_model()` | — |
| Multi-output | `outputs = list(...)`, `loss =` per head | manual |
| Attention | `ggml_layer_attention()`, `ggml_attention()` | `ag_multihead_attention()` |
| Compile | `ggml_compile()` | manual optimizer |
| Train | `ggml_fit()` | `with_grad_tape` + `backward` + `opt$step` |
| Evaluate | `ggml_evaluate()` | manual loop |
| Predict | `ggml_predict()` | manual forward |
| Save / load | `ggml_save_model()` / `ggml_load_model()` | — |
| Callbacks | `ggml_callback_early_stopping()`, `ggml_schedule_*` | `lr_scheduler_*`, `clip_grad_norm` |
| Device | automatic | `ag_device()` / `ag_dtype()` |
