tests/testthat/test-onnx-nn.R

In ggmlR: 'GGML' Tensor Operations for Machine Learning

# Tests for ONNX neural network ops:
# Conv, BatchNorm, LayerNorm, MatMul, Gemm, Gather

run_onnx <- function(path, inputs, device = "cpu") {
  m <- onnx_load(path, device = device)
  res <- onnx_run(m, inputs)
  res[[1]]
}

# ── MatMul ───────────────────────────────────────────────────────

test_that("ONNX MatMul works for 2D", {
  path <- .onnx_make_matmul(M = 2L, K = 3L, N = 2L)
  a <- c(1, 2, 3, 4, 5, 6)
  b <- c(1, 2, 3, 4, 5, 6)
  result <- run_onnx(path, list(A = a, B = b))
  expected <- c(22, 28, 49, 64)
  expect_equal(as.numeric(result), expected, tolerance = 1e-3)
})

# ── LayerNormalization ───────────────────────────────────────────

test_that("ONNX LayerNormalization works (1D)", {
  path <- .onnx_make_layer_norm(c(4L), eps = 1e-5)
  x <- c(1, 2, 3, 4)
  result <- run_onnx(path, list(X = x))
  r <- as.numeric(result)
  expect_equal(length(r), 4)
  expect_true(abs(mean(r)) < 0.01)
  pop_sd <- sqrt(mean((x - mean(x))^2))
  expected <- (x - mean(x)) / pop_sd
  expect_equal(r, expected, tolerance = 0.01)
})

# ── Conv 2D ────────────────────────────────────────────────────

test_that("ONNX Conv 2D basic works", {
  inp <- .onnx_value_info("X", 1L, c(1L, 1L, 3L, 3L))
  outp <- .onnx_value_info("Y", 1L, c(1L, 1L, 2L, 2L))
  w_data <- rep(1.0, 4)
  w_raw <- unlist(lapply(w_data, .float_bytes))
  w_t  <- .onnx_tensor("W", c(1L, 1L, 2L, 2L), 1L, w_raw)
  w_vi <- .onnx_value_info("W", 1L, c(1L, 1L, 2L, 2L))
  node <- .onnx_node("Conv", c("X", "W"), "Y")
  graph <- .onnx_graph("test", list(node),
                        list(inp, w_vi), list(outp), list(w_t))
  path <- tempfile(fileext = ".onnx")
  writeBin(.onnx_model(graph), path)

  x <- c(1, 2, 3, 4, 5, 6, 7, 8, 9)
  result <- run_onnx(path, list(X = x))
  expect_equal(as.numeric(result), c(12, 16, 24, 28), tolerance = 1e-3)
})

test_that("ONNX Conv 2D with bias works", {
  inp <- .onnx_value_info("X", 1L, c(1L, 1L, 3L, 3L))
  outp <- .onnx_value_info("Y", 1L, c(1L, 1L, 2L, 2L))
  w_data <- rep(1.0, 4)
  w_raw <- unlist(lapply(w_data, .float_bytes))
  w_t  <- .onnx_tensor("W", c(1L, 1L, 2L, 2L), 1L, w_raw)
  w_vi <- .onnx_value_info("W", 1L, c(1L, 1L, 2L, 2L))
  b_raw <- .float_bytes(10.0)
  b_t  <- .onnx_tensor("B", c(1L), 1L, b_raw)
  b_vi <- .onnx_value_info("B", 1L, c(1L))
  node <- .onnx_node("Conv", c("X", "W", "B"), "Y")
  graph <- .onnx_graph("test", list(node),
                        list(inp, w_vi, b_vi), list(outp),
                        list(w_t, b_t))
  path <- tempfile(fileext = ".onnx")
  writeBin(.onnx_model(graph), path)

  x <- c(1, 2, 3, 4, 5, 6, 7, 8, 9)
  result <- run_onnx(path, list(X = x))
  expect_equal(as.numeric(result), c(12, 16, 24, 28) + 10, tolerance = 1e-3)
})

# ── BatchNormalization ─────────────────────────────────────────

test_that("ONNX BatchNormalization 2D works", {
  inp <- .onnx_value_info("X", 1L, c(2L, 4L))
  outp <- .onnx_value_info("Y", 1L, c(2L, 4L))
  ones  <- unlist(lapply(rep(1.0, 4), .float_bytes))
  zeros <- unlist(lapply(rep(0.0, 4), .float_bytes))
  scale_t <- .onnx_tensor("scale", c(4L), 1L, ones)
  bias_t  <- .onnx_tensor("bias",  c(4L), 1L, zeros)
  mean_t  <- .onnx_tensor("mean",  c(4L), 1L, zeros)
  var_t   <- .onnx_tensor("var",   c(4L), 1L, ones)
  scale_vi <- .onnx_value_info("scale", 1L, c(4L))
  bias_vi  <- .onnx_value_info("bias",  1L, c(4L))
  mean_vi  <- .onnx_value_info("mean",  1L, c(4L))
  var_vi   <- .onnx_value_info("var",   1L, c(4L))
  node <- .onnx_node("BatchNormalization",
                      c("X", "scale", "bias", "mean", "var"), "Y")
  graph <- .onnx_graph("test", list(node),
                        list(inp, scale_vi, bias_vi, mean_vi, var_vi),
                        list(outp),
                        list(scale_t, bias_t, mean_t, var_t))
  path <- tempfile(fileext = ".onnx")
  writeBin(.onnx_model(graph), path)

  x <- c(1, 2, 3, 4, 5, 6, 7, 8)
  result <- run_onnx(path, list(X = x))
  expect_equal(as.numeric(result), x, tolerance = 1e-3)
})

test_that("ONNX BatchNormalization 4D works", {
  inp <- .onnx_value_info("X", 1L, c(1L, 2L, 2L, 2L))
  outp <- .onnx_value_info("Y", 1L, c(1L, 2L, 2L, 2L))
  scale_raw <- unlist(lapply(c(2.0, 1.0), .float_bytes))
  bias_raw  <- unlist(lapply(c(1.0, 0.0), .float_bytes))
  zeros2    <- unlist(lapply(rep(0.0, 2), .float_bytes))
  ones2     <- unlist(lapply(rep(1.0, 2), .float_bytes))
  scale_t <- .onnx_tensor("scale", c(2L), 1L, scale_raw)
  bias_t  <- .onnx_tensor("bias",  c(2L), 1L, bias_raw)
  mean_t  <- .onnx_tensor("mean",  c(2L), 1L, zeros2)
  var_t   <- .onnx_tensor("var",   c(2L), 1L, ones2)
  scale_vi <- .onnx_value_info("scale", 1L, c(2L))
  bias_vi  <- .onnx_value_info("bias",  1L, c(2L))
  mean_vi  <- .onnx_value_info("mean",  1L, c(2L))
  var_vi   <- .onnx_value_info("var",   1L, c(2L))
  node <- .onnx_node("BatchNormalization",
                      c("X", "scale", "bias", "mean", "var"), "Y")
  graph <- .onnx_graph("test", list(node),
                        list(inp, scale_vi, bias_vi, mean_vi, var_vi),
                        list(outp),
                        list(scale_t, bias_t, mean_t, var_t))
  path <- tempfile(fileext = ".onnx")
  writeBin(.onnx_model(graph), path)

  x <- c(1, 2, 3, 4, 5, 6, 7, 8)
  result <- run_onnx(path, list(X = x))
  expect_equal(as.numeric(result), c(3, 5, 7, 9, 5, 6, 7, 8), tolerance = 1e-3)
})

# ── Gather ─────────────────────────────────────────────────────

test_that("ONNX Gather works (embedding lookup)", {
  inp_idx <- .onnx_value_info("I", 7L, c(2L))
  w_data <- c(0.1, 0.2, 0.3, 0.4, 0.5, 0.6,
              0.7, 0.8, 0.9, 1.0, 1.1, 1.2)
  w_raw <- unlist(lapply(w_data, .float_bytes))
  w_t  <- .onnx_tensor("W", c(4L, 3L), 1L, w_raw)
  w_vi <- .onnx_value_info("W", 1L, c(4L, 3L))
  outp <- .onnx_value_info("Y", 1L, c(2L, 3L))
  node <- .onnx_node("Gather", c("W", "I"), "Y",
                      attrs = list(.onnx_attr_int("axis", 0L)))
  graph <- .onnx_graph("test", list(node),
                        list(inp_idx, w_vi), list(outp), list(w_t))
  path <- tempfile(fileext = ".onnx")
  writeBin(.onnx_model(graph), path)

  indices <- c(0, 2)
  result <- run_onnx(path, list(I = indices))
  expect_equal(as.numeric(result), c(0.1, 0.2, 0.3, 0.7, 0.8, 0.9), tolerance = 1e-3)
})