 /// empty or the topmost device is `nil`, that allows TensorFlow to place operations on any device
 /// that it sees fit.
 @usableFromInline
-class _ThreadLocalState {
-    var deviceScopes: [String?] = []
-
-    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
-    }()
+struct DeviceScopes {
+    var deviceStack: [String?] = []
 
     var _currentDevice: String? {
-        return deviceScopes.last ?? nil
+        return deviceStack.last ?? nil
     }
 
     @usableFromInline
-    func pushDevice(_ device: String?) {
-        deviceScopes.append(device)
+    mutating func pushDevice(_ device: String?) {
+        deviceStack.append(device)
     }
 
     @usableFromInline
-    func popDevice() {
-        internalConsistencyCheck(deviceScopes.popLast() != nil)
-    }
-
-    @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
+    mutating func popDevice() {
+        internalConsistencyCheck(deviceStack.popLast() != nil)
     }
 }
 

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

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

    }
}

/// 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.
