Persist classes from parents in subclasses

R/nn.R

@@ -1,6 +1,17 @@
 #' @include utils-data.R
 NULL
 
+get_inherited_classes <- function(inherit) {
+  inherit_class <- inherit$public_fields$.classes
+  # Filter out classes that we eventually add in our normal flow.
+  inherit_class <- inherit_class[inherit_class != "nn_Module"]
+  inherit_class <- inherit_class[inherit_class != "R6ClassGenerator"]
+  inherit_class <- inherit_class[inherit_class != "nn_module_generator"]
+  inherit_class <- inherit_class[inherit_class != "nn_module"]
+  inherit_class <- inherit_class[!duplicated(inherit_class, fromLast = TRUE)]
+  inherit_class
+}
+
 nn_Module <- R6::R6Class(
   classname = "nn_Module",
   lock_objects = FALSE,
@@ -443,8 +454,9 @@ nn_module <- function(classname = NULL, inherit = nn_Module, ...,
 
   e <- new.env(parent = parent_env)
   e$inherit <- inherit
-
-  classes <- c(classname, "nn_module")
+
+  inherit_class <- get_inherited_classes(inherit)
+  classes <- c(classname, inherit_class, "nn_module")
 
   Module <- R6::R6Class(
     classname = classname,
@@ -765,4 +777,3 @@ get_parameter_count <- function(self) {
   pars <- sapply(self$parameters, function(x) prod(x$shape))
   sum(pars)
 }
-

tests/testthat/test-nn.R

@@ -576,3 +576,30 @@ test_that("we can subset `nn_sequential`", {
   expect_true(inherits(y[[2]], "nn_relu6"))
 
 })
+
+test_that("classes are inherited correctly", {
+
+  nn <- nn_module(
+    classname = "hello",
+    inherit = nn_linear
+  )
+
+  nn2 <- nn_module(
+    classname = "goodbye",
+    inherit = nn
+  )
+
+  expect_equal(
+    class(nn), c("hello", "nn_linear", "nn_module", "nn_module_generator")
+  )
+
+  expect_equal(
+    class(nn2), c("goodbye", "hello", "nn_linear", "nn_module", "nn_module_generator")
+  )
+
+  n <- nn(10, 10)
+  expect_equal(class(n), c("hello", "nn_linear", "nn_module"))
+  n2 <- nn2(10, 10)
+  expect_equal(class(n2), c("goodbye", "hello", "nn_linear", "nn_module"))
+
+})