Overrides on ParamLessModule

Author: NiklasGustafsson

src/TorchSharp/NN/ParamLessModule.cs

@@ -9,11 +9,14 @@ namespace TorchSharp
 
     namespace Modules
     {
+        public interface IParameterLessModule {
+
+        }
         /// <summary>
         /// Base class for all modules that do not have any tensor parameters or buffers, and
         /// for which the `_to()` implementation can therefore be simplified.
         /// </summary>
-        public abstract class ParamLessModule<T1, T2> : nn.Module<T1, T2>
+        public abstract class ParamLessModule<T1, T2> : nn.Module<T1, T2>, IParameterLessModule
         {
             protected ParamLessModule(string name) : base(name) { }
 
@@ -26,13 +29,30 @@ protected ParamLessModule(IntPtr handle, IntPtr boxedHandle) : base(handle, boxe
             protected internal override nn.Module _to(DeviceType deviceType, int deviceIndex = -1) => this;
 
             protected internal override nn.Module _to(ScalarType dtype) => this;
+
+            public override void register_buffer(string name, Tensor tensor, bool persistent = true)
+            {
+                throw new InvalidOperationException($"Cannot register a buffer on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_parameter(string name, Parameter param)
+            {
+                throw new InvalidOperationException($"Cannot register a parameter on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_module(string name, nn.Module submodule)
+            {
+                if (submodule is not IParameterLessModule)
+                    throw new InvalidOperationException($"Submodules of a parameter-less module must also be parameter-less.");
+                base.register_module(name, submodule);
+            }
         }
 
         /// <summary>
         /// Base class for all modules that do not have any tensor parameters or buffers, and
         /// for which the `_to()` implementation can therefore be simplified.
         /// </summary>
-        public abstract class ParamLessModule<T1, T2, T3> : nn.Module<T1, T2, T3>
+        public abstract class ParamLessModule<T1, T2, T3> : nn.Module<T1, T2, T3>, IParameterLessModule
         {
             protected ParamLessModule(string name) : base(name) { }
 
@@ -45,13 +65,30 @@ protected ParamLessModule(IntPtr handle, IntPtr boxedHandle) : base(handle, boxe
             protected internal override nn.Module _to(DeviceType deviceType, int deviceIndex = -1) => this;
 
             protected internal override nn.Module _to(ScalarType dtype) => this;
+
+            public override void register_buffer(string name, Tensor tensor, bool persistent = true)
+            {
+                throw new InvalidOperationException($"Cannot register a buffer on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_parameter(string name, Parameter param)
+            {
+                throw new InvalidOperationException($"Cannot register a parameter on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_module(string name, nn.Module submodule)
+            {
+                if (submodule is not IParameterLessModule)
+                    throw new InvalidOperationException($"Submodules of a parameter-less module must also be parameter-less.");
+                base.register_module(name, submodule);
+            }
         }
 
         /// <summary>
         /// Base class for all modules that do not have any tensor parameters or buffers, and
         /// for which the `_to()` implementation can therefore be simplified.
         /// </summary>
-        public abstract class ParamLessModule<T1, T2, T3, T4> : nn.Module<T1, T2, T3, T4>
+        public abstract class ParamLessModule<T1, T2, T3, T4> : nn.Module<T1, T2, T3, T4>, IParameterLessModule
         {
             protected ParamLessModule(string name) : base(name) { }
 
@@ -64,6 +101,23 @@ protected ParamLessModule(IntPtr handle, IntPtr boxedHandle) : base(handle, boxe
             protected internal override nn.Module _to(DeviceType deviceType, int deviceIndex = -1) => this;
 
             protected internal override nn.Module _to(ScalarType dtype) => this;
+
+            public override void register_buffer(string name, Tensor tensor, bool persistent = true)
+            {
+                throw new InvalidOperationException($"Cannot register a buffer on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_parameter(string name, Parameter param)
+            {
+                throw new InvalidOperationException($"Cannot register a parameter on a module that is declared 'parameter-less.'");
+            }
+
+            public override void register_module(string name, nn.Module submodule)
+            {
+                if (submodule is not IParameterLessModule)
+                    throw new InvalidOperationException($"Submodules of a parameter-less module must also be parameter-less.");
+                base.register_module(name, submodule);
+            }
         }
     }
 }

test/TorchSharpTest/NN.cs

@@ -6633,5 +6633,37 @@ public void TestModulePostHooks()
             lin1.call(input);
             Assert.Equal(1, counter);
         }
+
+        [Fact]
+        public void TestCustomParameterLessModule() 
+        {
+            var cnp = new CustomNoParameters("test");
+
+            // Should not throw
+            cnp.register_module("sub", new CustomNoParameters("test"));
+
+            Assert.True(cnp.named_modules().Count() > 0);
+            Assert.Equal("sub", cnp.named_modules().First().name);
+
+            Assert.Throws<InvalidOperationException>(() => cnp.register_module("test", torch.nn.Linear(10,10, true)));
+            Assert.Throws<InvalidOperationException>(() => cnp.register_buffer("test", torch.rand(10)));
+            Assert.Throws<InvalidOperationException>(() => cnp.register_parameter("test", new Parameter(torch.rand(10))));
+        }
+
+        class CustomNoParameters : ParamLessModule<Tensor, Tensor>
+        {
+            public CustomNoParameters(string name) : base(name)
+            {
+            }
+
+            public CustomNoParameters(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle)
+            {
+            }
+
+            public override Tensor forward(Tensor input)
+            {
+                throw new NotImplementedException();
+            }
+        }
     }
 }