Module hooks don't depend on the forward() signature.

Add CopyTo/From overloads to TensorAccessor for Span<T>.

Author: NiklasGustafsson

src/TorchSharp/NN/Module.cs

@@ -1230,7 +1230,7 @@ public interface IModule<T1, T2, T3, T4, T5, T6, TResult>
             /// <summary>
             /// Represents a module that accepts 'hook' to the module logic.
             /// </summary>
-            public class HookableModule<TModule,TPreHook,TPostHook> : Module
+            public class HookableModule<TPreHook,TPostHook> : Module
             {
                 protected HookableModule(string name) : base(name) { }
 
@@ -1251,14 +1251,14 @@ public HookRemover register_forward_pre_hook(TPreHook hook)
                     return new HookRemover(this, key);
                 }
 
-                public HookRemover register_forward_hook(Action<TModule> hook)
+                public HookRemover register_forward_hook(Action<Module> hook)
                 {
                     var key = Guid.NewGuid().ToString();
                     module_post_hooks.Add(key, hook);
                     return new HookRemover(this, key);
                 }
 
-                public HookRemover register_forward_pre_hook(Action<TModule> hook)
+                public HookRemover register_forward_pre_hook(Action<Module> hook)
                 {
                     var key = Guid.NewGuid().ToString();
                     module_pre_hooks.Add(key, hook);
@@ -1276,16 +1276,16 @@ private void remove(string key)
                 protected Dictionary<string, TPreHook> pre_hooks = new Dictionary<string, TPreHook>();
                 protected Dictionary<string, TPostHook> post_hooks = new Dictionary<string, TPostHook>();
 
-                protected Dictionary<string, Action<TModule>> module_pre_hooks = new Dictionary<string, Action<TModule>>();
-                protected Dictionary<string, Action<TModule>> module_post_hooks = new Dictionary<string, Action<TModule>>();
+                protected Dictionary<string, Action<Module>> module_pre_hooks = new Dictionary<string, Action<Module>>();
+                protected Dictionary<string, Action<Module>> module_post_hooks = new Dictionary<string, Action<Module>>();
 
                 /// <summary>
                 /// Used to remove a specific hook, following the PyTorch API design.
                 /// </summary>
                 /// <remarks>The name and namespace of this class is not the same as in PyTorch, but serves the same purpose.</remarks>
                 public class HookRemover
                 {
-                    public HookRemover(HookableModule<TModule,TPreHook, TPostHook> module, string key)
+                    public HookRemover(HookableModule<TPreHook, TPostHook> module, string key)
                     {
                         this.module = module;
                         this.key = key;
@@ -1296,7 +1296,7 @@ public void remove()
                         module.remove(key);
                     }
 
-                    private HookableModule<TModule,TPreHook, TPostHook> module;
+                    private HookableModule<TPreHook, TPostHook> module;
                     private string key;
                 }
             }
@@ -1306,7 +1306,7 @@ public void remove()
             /// </summary>
             /// <typeparam name="T">The argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T, TResult> : HookableModule<Module<T, TResult>, Func<Module<T,TResult>, T, T>, Func<Module<T, TResult>, T, TResult, TResult>>, IModule<T, TResult>
+            public abstract class Module<T, TResult> : HookableModule<Func<Module<T,TResult>, T, T>, Func<Module<T, TResult>, T, TResult, TResult>>, IModule<T, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1360,7 +1360,7 @@ public TResult call(T input)
             /// <typeparam name="T1">The first argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T2">The second argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, TResult> : HookableModule<Module<T1, T2, TResult>, Func<Module<T1, T2, TResult>, T1, T2, (T1, T2)?>, Func<Module<T1, T2, TResult>, T1, T2, TResult, TResult>>, IModule<T1, T2, TResult>
+            public abstract class Module<T1, T2, TResult> : HookableModule<Func<Module<T1, T2, TResult>, T1, T2, (T1, T2)?>, Func<Module<T1, T2, TResult>, T1, T2, TResult, TResult>>, IModule<T1, T2, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1417,7 +1417,7 @@ public TResult call(T1 input1, T2 input2)
             /// <typeparam name="T2">The second argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T3">The third argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, TResult> : HookableModule<Module<T1, T2, T3, TResult>,Func<Module<T1, T2, T3, TResult>, T1, T2, T3, (T1, T2, T3)?>, Func<Module<T1, T2, T3, TResult>, T1, T2, T3, TResult, TResult>>, IModule<T1, T2, T3, TResult>
+            public abstract class Module<T1, T2, T3, TResult> : HookableModule<Func<Module<T1, T2, T3, TResult>, T1, T2, T3, (T1, T2, T3)?>, Func<Module<T1, T2, T3, TResult>, T1, T2, T3, TResult, TResult>>, IModule<T1, T2, T3, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1476,7 +1476,7 @@ public TResult call(T1 input1, T2 input2, T3 input3)
             /// <typeparam name="T3">The third argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T4">The fourth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, TResult> : HookableModule<Module<T1, T2, T3, T4, TResult>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, (T1, T2, T3, T4)?>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, TResult, TResult>>, IModule<T1, T2, T3, T4, TResult>
+            public abstract class Module<T1, T2, T3, T4, TResult> : HookableModule<Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, (T1, T2, T3, T4)?>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, TResult, TResult>>, IModule<T1, T2, T3, T4, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1537,7 +1537,7 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4)
             /// <typeparam name="T4">The fourth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T5">The fifth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, T5,TResult> : HookableModule<Module<T1, T2, T3, T4, T5,TResult>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, (T1, T2, T3, T4, T5)?>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, TResult>
+            public abstract class Module<T1, T2, T3, T4, T5,TResult> : HookableModule<Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, (T1, T2, T3, T4, T5)?>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1600,7 +1600,7 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5)
             /// <typeparam name="T5">The fifth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T6">The sixth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, T5, T6, TResult> : HookableModule<Module<T1, T2, T3, T4, T5, T6, TResult>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, (T1, T2, T3, T4, T5, T6)?>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, T6, TResult>
+            public abstract class Module<T1, T2, T3, T4, T5, T6, TResult> : HookableModule<Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, (T1, T2, T3, T4, T5, T6)?>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, T6, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }

src/TorchSharp/Utils/TensorAccessor.cs

@@ -241,6 +241,16 @@ public void CopyTo(T[] array, int arrayIndex = 0, long tensorIndex = 0)
             }
         }
 
+        public void CopyTo(Span<T> array, int arrayIndex = 0, long tensorIndex = 0)
+        {
+            int idx = arrayIndex;
+            foreach (int offset in GetSubsequentIndices(tensorIndex)) {
+                if (idx >= array.Length) break;
+                unsafe { array[idx] = ((T*)_tensor_data_ptr)[offset]; }
+                idx += 1;
+            }
+        }
+
         public void CopyFrom(T[] array, int arrayIndex = 0, long tensorIndex = 0)
         {
             int idx = arrayIndex;
@@ -251,6 +261,16 @@ public void CopyFrom(T[] array, int arrayIndex = 0, long tensorIndex = 0)
             }
         }
 
+        public void CopyFrom(ReadOnlySpan<T> array, int arrayIndex = 0, long tensorIndex = 0)
+        {
+            int idx = arrayIndex;
+            foreach (int offset in GetSubsequentIndices(tensorIndex)) {
+                if (idx >= array.Length) break;
+                unsafe { ((T*)_tensor_data_ptr)[offset] = array[idx]; }
+                idx += 1;
+            }
+        }
+
         /// <summary>
         /// Translates a linear index within the span represented by the accessor to a linear index
         /// used by the underlying tensor. The two should only be different if the tensor is a view

test/TorchSharpTest/NN.cs

@@ -6591,7 +6591,7 @@ public void TestModulePreHooksGeneric()
             var input = torch.randn(32, 100, 100);
             var counter = 0;
 
-            var pre_hook = (Module<Tensor, Tensor> m) => { counter += 1;};
+            var pre_hook = (Module m) => { counter += 1;};
 
             var handle = lin1.register_forward_pre_hook(pre_hook);
 
@@ -6611,7 +6611,7 @@ public void TestModulePostHooksGeneric()
             var input = torch.randn(32, 100, 100);
             var counter = 0;
 
-            var hook = (Module<Tensor, Tensor> m) => { counter += 1;};
+            var hook = (Module m) => { counter += 1;};
 
             var handle = lin1.register_forward_hook(hook);