         guard deviceNames.contains(name) else {
             fatalError("Device \(name) not found")
         }
-        _ThreadLocalState.local.pushDevice(name)
+        _ThreadLocalState.local.deviceScopes.pushDevice(name)
         let result = try body()
-        _ThreadLocalState.local.popDevice()
+        _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.pushDevice(nil)
+        _ThreadLocalState.local.deviceScopes.pushDevice(nil)
         let result = try body()
-        _ThreadLocalState.local.popDevice()
+        _ThreadLocalState.local.deviceScopes.popDevice()
         return result
     }
 }

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

    }
}

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