Track the device name in LazyTensorOperation

Sources/TensorFlow/Core/LazyTensorOperation.swift

@@ -162,6 +162,7 @@ class LazyTensorOperation: TensorOperation {
     let outputCount: Int
     var inputs: [Input]
     var attributes: [String: Attribute]
+    var deviceName: String?
     var outputs: [TFETensorHandle]?
     var id: String?
 
@@ -203,6 +204,7 @@ class LazyTensorOperation: TensorOperation {
         self.name = name
         self.inputs = []
         self.attributes = [:]
+        self.deviceName = _ExecutionContext.global.currentDeviceName
         self.outputCount = outputCount
         self.outputs = nil
         self.id = id

Tests/TensorFlowTests/LazyTensorOperationTests.swift

@@ -17,6 +17,7 @@ import XCTest
 import CTensorFlow
 
 final class LazyTensorOperationTests: XCTestCase {
+
     func testNoInput() {
         let placeholder = LazyTensorOperation(
             _id: "V", name: "Placeholder", outputCount: 1)
@@ -204,6 +205,15 @@ final class LazyTensorOperationTests: XCTestCase {
         XCTAssertEqual(op0.description, "%0 = Nop[fn: TFFunction(ExampleFunction)]()")
     }
 
+    func testDeviceTracking() {
+        let op0 = LazyTensorOperation(_id: "0", name: "Nop", outputCount: 1)
+        XCTAssertEqual(op0.deviceName, nil)
+        withDevice(named: "/job:localhost/replica:0/task:0/device:CPU:0") {
+            let op1 = LazyTensorOperation(_id: "0", name: "Nop", outputCount: 1)
+            XCTAssertEqual(op1.deviceName ?? "", "/job:localhost/replica:0/task:0/device:CPU:0")
+        }
+    }
+
     static var allTests = [
         ("testNoInput", testNoInput),
         ("testSingleInput", testSingleInput),
@@ -223,6 +233,8 @@ final class LazyTensorOperationTests: XCTestCase {
             testOptionalTensorShapeArrayAttribute),
         ("testArrayAttributes", testArrayAttributes),
         ("testMultipleAttributes", testMultipleAttributes),
-        ("testFunctionAttribute", testFunctionAttribute)
+        ("testFunctionAttribute", testFunctionAttribute),
+        ("testDeviceTracking", testDeviceTracking)
+
     ]
 }