Preparing spec_helpers for functional support

Author: davidrsch

R/create_keras_spec_helpers.R

@@ -1,20 +1,36 @@
 #' Discover and Collect Model Specification Arguments
 #'
-#' Introspects the provided layer block functions to generate a list of
-#' arguments for the new model specification. This includes arguments for
-#' block repetition (`num_*`), block-specific hyperparameters (`block_*`),
-#' and global training parameters.
+#' @description
+#' This internal helper introspects the user-provided `layer_blocks` functions
+#' to generate a complete list of arguments for the new model specification.
+#' The logic for discovering arguments differs for sequential and functional models.
+#'
+#' @details
+#' For **sequential models** (`functional = FALSE`):
+#' - It creates `num_{block_name}` arguments to control block repetition.
+#' - It inspects the arguments of each block function, skipping the first
+#'   (assumed to be the `model` object), to find tunable hyperparameters.
+#'
+#' For **functional models** (`functional = TRUE`):
+#' - It does **not** create `num_{block_name}` arguments.
+#' - It inspects the arguments of each block function. Arguments whose names
+#'   match other block names are considered graph connections (inputs) and are
+#'   ignored. The remaining arguments are treated as tunable hyperparameters.
+#'
+#' In both cases, it also adds global training parameters (like `epochs`) and
+#' filters out special engine-supplied arguments (`input_shape`, `num_classes`).
 #'
 #' @param layer_blocks A named list of functions defining Keras layer blocks.
+#' @param functional A logical. If `TRUE`, uses discovery logic for the
+#'   Functional API. If `FALSE`, uses logic for the Sequential API.
 #' @param global_args A character vector of global arguments to add to the
 #'   specification (e.g., "epochs").
 #' @return A list containing two elements:
-#'   - `all_args`: A named list of arguments for the new function signature,
-#'     initialized with `rlang::zap()`.
-#'   - `parsnip_names`: A character vector of all argument names for `parsnip`.
+#'
 #' @noRd
 collect_spec_args <- function(
   layer_blocks,
+  functional,
   global_args = c(
     "epochs",
     "batch_size",
@@ -36,23 +52,39 @@ collect_spec_args <- function(
   all_args <- list()
   parsnip_names <- character()
 
+  block_names <- names(layer_blocks)
+
   # block repetition counts (e.g., num_dense)
-  for (block in names(layer_blocks)) {
-    num_name <- paste0("num_", block)
+  for (block_name in block_names) {
+    num_name <- paste0("num_", block_name)
     all_args[[num_name]] <- rlang::zap()
     parsnip_names <- c(parsnip_names, num_name)
   }
 
   # These args are passed by the fit engine, not set by the user in the spec
   engine_args <- c("input_shape", "num_classes")
-  # block-specific parameters (skip first 'model' formal)
-  for (block in names(layer_blocks)) {
-    fmls_to_process <- rlang::fn_fmls(layer_blocks[[block]])[-1]
-    # Filter out arguments that are provided by the fitting engine
-    for (arg in names(fmls_to_process[
-      !names(fmls_to_process) %in% engine_args
-    ])) {
-      full <- paste0(block, "_", arg)
+  # Discover block-specific hyperparameters
+  for (block_name in block_names) {
+    block_fmls <- rlang::fn_fmls(layer_blocks[[block_name]])
+
+    if (isTRUE(functional)) {
+      # For functional models, hyperparameters are arguments that are NOT
+      # names of other blocks (which are graph connections).
+      hyperparam_names <- setdiff(
+        names(block_fmls),
+        c(block_names, engine_args)
+      )
+    } else {
+      # For sequential models, hyperparameters are all args except the first
+      # ('model') and special engine args.
+      fmls_to_process <- if (length(block_fmls) > 0) block_fmls[-1] else list()
+      hyperparam_names <- names(fmls_to_process)[
+        !names(fmls_to_process) %in% engine_args
+      ]
+    }
+
+    for (arg in hyperparam_names) {
+      full <- paste0(block_name, "_", arg)
       all_args[[full]] <- rlang::zap()
       parsnip_names <- c(parsnip_names, full)
     }
@@ -69,8 +101,10 @@ collect_spec_args <- function(
 
 #' Internal Implementation for Creating Keras Specifications
 #'
-#' This is the core logic for both `create_keras_sequential_spec` and
-#' `create_keras_functional_spec`. It is not intended for direct use.
+#' @description
+#' This is the core implementation for both `create_keras_sequential_spec()` and
+#' `create_keras_functional_spec()`. It orchestrates the argument collection,
+#' function building, and `parsnip` registration steps.
 #'
 #' @inheritParams create_keras_sequential_spec
 #' @param functional A logical, if `TRUE`, registers the model to be fit with
@@ -85,7 +119,7 @@ create_keras_spec_impl <- function(
   functional,
   env
 ) {
-  args_info <- collect_spec_args(layer_blocks)
+  args_info <- collect_spec_args(layer_blocks, functional = functional)
   spec_fun <- build_spec_function(
     model_name,
     mode,