Make useLazyTensor a thread-local state. (#387)

Author: bgogul

Sources/TensorFlow/Core/LazyTensorTrace.swift

@@ -91,7 +91,7 @@ class LazyTensorTraceBuilder {
 
     /// Returns a trace obtained by tracing the given function.
     static func trace<In: TensorGroup, Out: TensorGroup>(_ fn: (In) -> Out) -> LazyTensorTrace {
-        precondition(_RuntimeConfig.useLazyTensor, "Lazy tensor is not enabled for tracing.")
+        precondition(_ThreadLocalState.useLazyTensor, "Lazy tensor is not enabled for tracing.")
 
         // Set up inputs for running `fn`.
         let inputOps = In._typeList.map { Self.makePlaceholder(dataType: $0) }

Sources/TensorFlow/Core/Runtime.swift

@@ -693,7 +693,7 @@ extension _ExecutionContext {
     static func makeOp(
         _ name: String, _ outputCount: Int
     ) -> TFTensorOperation {
-        return _RuntimeConfig.useLazyTensor
+        return _ThreadLocalState.useLazyTensor
             ? LazyTensorOperation(name, outputCount)
             : TFE_Op(name, outputCount)
     }
@@ -1201,6 +1201,20 @@ class _ThreadLocalState {
 
     var lazyTensorContext = LazyTensorContext()
 
+    static var useLazyTensor: Bool {
+        get {
+            _ThreadLocalState.local.lazyTensorEnabled ?? _RuntimeConfig.useLazyTensor
+        }
+        set {
+            _ThreadLocalState.local.lazyTensorEnabled = newValue
+        }
+    }
+
+    /// When true, use lazy evaluation. If this is not set, we should use the
+    /// value of `_RuntimeConfig.useLazyTensor` to determine if lazy evaluation
+    /// is enabled.
+    private var lazyTensorEnabled: Bool? = nil
+
     private static let key: pthread_key_t = {
         var key = pthread_key_t()
         pthread_key_create(&key) {