Fix the creation and replacement of a Tensor with a Parameter instance

Author: NiklasGustafsson

src/TorchSharp/DisposeScope.cs

@@ -224,6 +224,21 @@ public void Detach(IEnumerable<IDisposable> disposables)
             }
         }
 
+        /// <summary>
+        /// Replaces registration of one tensor with another.
+        /// </summary>
+        /// <param name="original">The original tensor, possibly registered under a dispose scope.</param>
+        /// <param name="replacement">The replacement tensor.</param>
+        internal static void ReplaceWith(torch.Tensor original, torch.Tensor replacement)
+        {
+            DisposeScope? scope = original.OwningDisposeScope;
+
+            if (scope != null && scope.Disposables.Remove(original)) {
+                original.OwningDisposeScope = null;
+                AddToOther(scope, replacement);
+            }
+        }
+
         public void Attach(IDisposable disposable)
         {
             _ = Attach((IEnumerable<IDisposable>)new[] { disposable });
@@ -369,10 +384,10 @@ public void MarkAsDisposed(IDisposable disposable)
         /// <returns></returns>
         public bool Contains(IDisposable disposable) => Disposables.Contains(disposable);
 
-        private bool AddToOther(DisposeScope scope, IDisposable disposable)
+        private static bool AddToOther(DisposeScope scope, IDisposable disposable)
         {
-            if (this._disposeScopeManager is null)
-                throw new ObjectDisposedException(this.GetType().FullName);
+            // if (this._disposeScopeManager is null)
+            //     throw new ObjectDisposedException(this.GetType().FullName);
 
             DisposeScope? oldScope;
             if (disposable is torch.Tensor t) {

src/TorchSharp/NN/Parameter.cs

@@ -24,13 +24,11 @@ public class Parameter : Tensor
             /// <param name="data">A tensor, which will become empty.</param>
             /// <param name="requires_grad"></param>
             public Parameter(Tensor data, bool requires_grad = true) :
-                base(data.with_requires_grad(requires_grad).MoveHandle())
+                base(data.with_requires_grad(requires_grad).MoveHandle(), false)
             {
                 var scope = data.OwningDisposeScope;
                 if (scope is not null) {
-                    this.OwningDisposeScope = scope;
-                    scope.Attach(this);
-                    scope.Detach(data);
+                    DisposeScope.ReplaceWith(data, this);
                 }
             }
 

src/TorchSharp/Tensor/Tensor.cs

@@ -34,12 +34,14 @@ public partial class Tensor : IDisposable
 
             internal DisposeScope? OwningDisposeScope { get; set; }
 
-            internal Tensor(IntPtr handle)
+            internal Tensor(IntPtr handle, bool register = true)
             {
                 this.handle = handle;
                 System.Threading.Interlocked.Increment(ref _totalCount);
                 _peakCount = Math.Max(_totalCount, _peakCount);
-                OwningDisposeScope = DisposeScopeManager.ThreadSingleton.RegisterOnCurrentDisposeScope(this);
+                if (register) {
+                    OwningDisposeScope = DisposeScopeManager.ThreadSingleton.RegisterOnCurrentDisposeScope(this);
+                }
             }
 
             /// <summary>
@@ -213,6 +215,11 @@ public IntPtr Handle {
                 }
             }
 
+            /// <summary>
+            /// Disassociates the native tensor handle from the managed Tensor.
+            /// </summary>
+            /// <remarks>Used to create a Parameter instance.</remarks>
+            /// <returns>The handle to the underlying native tensor.</returns>
             internal IntPtr MoveHandle()
             {
                 var h = handle;
@@ -2172,7 +2179,7 @@ public Tensor transpose_(long dim0, long dim1)
                 CheckForErrors();
                 return this;
             }
-            
+
             public Tensor threshold(Scalar threshold, Scalar value)
             {
                 var res = NativeMethods.THSTensor_threshold(Handle, threshold.Handle, value.Handle);
@@ -2717,7 +2724,7 @@ public Tensor softmax(long dim, ScalarType? dtype = null) =>
                 torch.special.softmax(this, dim, dtype);
 
 
-            public Tensor softplus(int beta = 1, int threshold = 20) => 
+            public Tensor softplus(int beta = 1, int threshold = 20) =>
                 softplus1(beta, threshold);
 
             private Tensor softplus1(Scalar beta, Scalar threshold)
@@ -2787,9 +2794,9 @@ public Tensor rrelu_(double lower = one_eighth, double upper = one_third)
             }
 
             public Tensor celu() => this.celu(1.0);
-            
+
             public Tensor celu_() => this.celu_(1.0);
-            
+
             public Tensor celu(Scalar alpha)
             {
                 var res = NativeMethods.THSTensor_celu(Handle, alpha.Handle);

test/TorchSharpTest/TestDisposeScopes.cs

@@ -192,7 +192,7 @@ public void DisposeScopesCanBeNestled()
             DisposeScopeManager.Statistics.Reset();
         }
 
-        [Fact(Skip = "https://github.com/dotnet/TorchSharp/issues/1397")]
+        [Fact]
         public void DisposeScopeWorksForTestTraining1()
         {
             DisposeScopeManager.Statistics.Reset();
@@ -207,7 +207,7 @@ public void DisposeScopeWorksForTestTraining1()
             DisposeScopeManager.Statistics.Reset();
         }
 
-        [Fact(Skip = "https://github.com/dotnet/TorchSharp/issues/1397")]
+        [Fact]
         public void DisposeScopeWorksForTestTrainingConv2d()
         {
             DisposeScopeManager.Statistics.Reset();