     }
 }
 
+/// A class to keep around thread local state.
+class _ThreadLocalState {
+    var deviceScopes = DeviceScopes()
+
+    private static let key: pthread_key_t = {
+        var key = pthread_key_t()
+        pthread_key_create(&key) {
+#if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
+            let _: AnyObject = Unmanaged.fromOpaque($0).takeRetainedValue()
+#else
+            let _: AnyObject = Unmanaged.fromOpaque($0!).takeRetainedValue()
+#endif
+        }
+        return key
+    }()
+
+    @usableFromInline
+    static var local: _ThreadLocalState {
+        if let state = pthread_getspecific(key) {
+            return Unmanaged.fromOpaque(state).takeUnretainedValue()
+        }
+        let state = _ThreadLocalState()
+        pthread_setspecific(key, Unmanaged.passRetained(state).toOpaque())
+        return state
+    }
+}
+
 /// Stack of devices that models nested calls to withDevice/withDefaultDevice. Devices are
 /// represented by their names in TensorFlow notation. See documentation for
 /// `withDevice(named:perform:)` to learn about device names.

    /// withDevice call on the call stack or the presence of an immediately enclosing
    /// `withDefaultDevice(perform)` call.
    var currentDeviceName: String? {
        return _ThreadLocalState.local._currentDevice
    }

    /// See documentation for the top-level `withDevice(_:_:perform)`.

        guard deviceNames.contains(name) else {
            fatalError("Device \(name) not found")
        }
        _ThreadLocalState.local.deviceScopes.pushDevice(name)
        let result = try body()
        _ThreadLocalState.local.deviceScopes.popDevice()
        return result
    }

    /// See documentation for the top-level `withDefaultDevice(perform)`.
    func withDefaultDevice<R>(perform body: () throws -> R) rethrows -> R {
        _ThreadLocalState.local.deviceScopes.pushDevice(nil)
        let result = try body()
        _ThreadLocalState.local.deviceScopes.popDevice()
        return result
    }
}
