From 9e17ab1aef2e4f4fdaf08064ce0f48dd50701659 Mon Sep 17 00:00:00 2001
From: alexpantyukhin <apantykhin@gmail.com>
Date: Fri, 17 Nov 2017 20:34:44 +0400
Subject: [PATCH] add DefaultGraphStack

---
 TensorFlowSharp/Tensorflow.cs | 34 +++++++++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

diff --git a/TensorFlowSharp/Tensorflow.cs b/TensorFlowSharp/Tensorflow.cs
index 0f361989..b2fceb18 100644
--- a/TensorFlowSharp/Tensorflow.cs
+++ b/TensorFlowSharp/Tensorflow.cs
@@ -108,6 +108,36 @@ internal static void CheckSize ()
 		}
 	}
 
+	internal class DefaultGraphStack
+	{
+		[ThreadStatic]
+		private static IList<TFGraph> stack;
+		private static Object lockObj;
+
+		internal TFGraph Instance => stack.Last();
+
+		internal static void SetGraph(TFGraph graph)
+		{
+			if (stack == null)
+			{
+				lock (lockObj)
+				{
+					if (stack == null)
+					{
+						stack = new List<TFGraph>();
+					}
+				}
+			}
+
+			stack.Add(graph);
+		}
+
+		internal static void RemoveGraph(TFGraph graph)
+		{
+			stack.Remove(graph);
+		}
+	}
+
 	/// <summary>
 	/// Base class for many TensorFlow data types that provides a common idiom to dispose and
 	/// release resources associated with the native data types.   Generally, you do not need to use this.
@@ -473,12 +503,13 @@ public partial class TFGraph : TFDisposable
 		/// <summary>
 		/// Initializes a new instance of the <see cref="T:TensorFlow.TFGraph"/> class.
 		/// </summary>
-		public TFGraph () : base (TF_NewGraph ())
+		public TFGraph () : this (TF_NewGraph ())
 		{
 		}
 
 		internal TFGraph (IntPtr handle) : base (handle)
 		{
+			DefaultGraphStack.SetGraph(this);
 		}
 
 		// extern void TF_DeleteGraph (TF_Graph *);
@@ -486,6 +517,7 @@ internal TFGraph (IntPtr handle) : base (handle)
 		static extern unsafe void TF_DeleteGraph (TF_Graph graph);
 		internal override void NativeDispose (IntPtr handle)
 		{
+			DefaultGraphStack.RemoveGraph(this);
 			TF_DeleteGraph (handle);
 		}