Putting LazyTensor components together (#294)

Extracted traces are evaluated and used for materializing lazy tensors.

Author: bgogul

Sources/TensorFlow/Core/LazyTensorOperation.swift

@@ -33,9 +33,8 @@ class LazyTensor: _AnyTensorHandle {
         switch handle {
         case .concrete(let h, _):
             return h
-        case .symbolic(_, _, _):
-            fatalError("TODO: to be send out in a separate PR.")
-            // return op.materialized(index: index)
+        case .symbolic(let op, let index, _):
+            return op.materialized(index: index)
         }
     }
 
@@ -725,3 +724,76 @@ extension LazyTensorOperation: CustomStringConvertible {
         return desc
     }
 }
+
+extension LazyTensorOperation {
+    /// Returns the materialized value at the given output `index`.
+    func materialized(index: Int) -> TFETensorHandle {
+        precondition(index < outputCount)
+        return materialized()[index]
+    }
+
+    /// Materializes all the outputs.
+    func materialized() -> [TFETensorHandle] {
+        // Return materialized outputs if any.
+        if let outputs = outputs { return outputs }
+
+        materializeLiveTensors()
+
+        // Our outputs should have been updated by now. Otherwise,
+        // something terrible happened!
+        precondition(outputs != nil, "Materialization failed!")
+        return outputs!
+    }
+
+    /// Converts symbolic tensor inputs to concrete inputs if the
+    /// associated `LazyTensorOperation` has been materialized.
+    private func maybeMaterializeInputs() {
+        /// If `lazyTensor` is symbolic and the associated `LazyTensorOperation`
+        /// has been materialized, return the corresponding concrete `LazyTensor`.
+        /// Otherwise, return `lazyTensor` untouched.
+        func materializedAsNeeded(lazyTensor: LazyTensor) -> LazyTensor {
+            let handle = lazyTensor.handle
+            if case let LazyTensor.Handle.symbolic(lazyOp, index, _) = handle,
+                let outputs = lazyOp.outputs {
+                return LazyTensor(_materialized: outputs[index])
+            }
+            return lazyTensor
+        }
+
+        /// Returns an input that is rewritten such that all symbolic values
+        /// that have been materialized have been replaced by the corresponding
+        /// concerete inputs. If no symbolic values have been materialized or if
+        /// there are no symbolic values, return the `input` untouched.
+        func materializedAsNeeded(input: Input) -> Input {
+            switch input {
+            case .single(let h):
+                return .single(materializedAsNeeded(lazyTensor: h))
+            case .list(let elements):
+                return .list(elements.map { materializedAsNeeded(lazyTensor: $0) })
+            }
+        }
+        inputs = inputs.map { materializedAsNeeded(input: $0) }
+    }
+
+    private func materializeLiveTensors() {
+        let lazyTrace = LazyTensorTrace(self)
+        debugLog("Extracted trace:\n\(lazyTrace)")
+
+        let function = TFFunction(trace: lazyTrace)
+        debugLog("Generated TFFunction:\n\(function)")
+
+        let allOutputs = function.execute(lazyTrace.inputValues)
+
+        // Slice up the outputs to various lazy tensors
+        var start = 0
+        for lazyOp in lazyTrace.originalOutputs {
+            let end = start + lazyOp.outputCount
+            lazyOp.outputs = Array(allOutputs[start..<end])
+            start = end
+        }
+
+        // On all the live operations rewrite the inputs so that we drop references
+        // to the LazyTensorOperations.
+        LazyTensor.forEachOperation { $0.maybeMaterializeInputs() }
+    }
+}

Tests/TensorFlowTests/LazyTensorEvaluationTests.swift

@@ -0,0 +1,105 @@
+// Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import XCTest
+
+@testable import TensorFlow
+import CTensorFlow
+
+final class LazyTensorEvaluationTests: XCTestCase {
+    override class func setUp() {
+        super.setUp()
+        _RuntimeConfig.useLazyTensor = true
+    }
+
+    override class func tearDown() {
+        super.tearDown()
+        _RuntimeConfig.useLazyTensor = false
+    }
+
+    func testSimpleOperations() {
+        let a = Tensor<Float>(10.0)
+        let b = Tensor<Float>(2.0)
+        let c = Tensor<Float>(3.0)
+        let w = a + b * c
+
+        XCTAssertFalse(isMaterialized(w))
+        XCTAssertEqual(w.scalarized(), 16.0)
+        XCTAssertTrue(isMaterialized(w))
+    }
+
+    func testMultipleMaterializations() {
+        let a = Tensor<Float>(10.0)
+        let b = Tensor<Float>(2.0)
+        let c = Tensor<Float>(3.0)
+        let x = a + b + c
+        let y = x * c
+        let z = y / (x - c)
+
+        // Materialize y first
+        XCTAssertFalse(isMaterialized(x))
+        XCTAssertFalse(isMaterialized(y))
+        XCTAssertFalse(isMaterialized(z))
+        XCTAssertEqual(y.scalarized(), 45.0)
+
+        // x and y are materialized, but not z.
+        XCTAssertTrue(isMaterialized(x))
+        XCTAssertTrue(isMaterialized(y))
+        XCTAssertFalse(isMaterialized(z))
+
+        XCTAssertEqual(z.scalarized(), 3.75)
+        XCTAssertTrue(isMaterialized(z))
+    }
+
+    func testSimpleControlFlow() {
+        let a = Tensor<Float>(5.0)
+        let addOrMul = { (useAdd: Bool, a: Tensor<Float>) in
+            useAdd ? (a + a) : (a * a)
+        }
+        let add = addOrMul(/*useAdd:*/true, a)
+        XCTAssertFalse(isMaterialized(add))
+        XCTAssertEqual(add.scalarized(), 10.0);
+        XCTAssertTrue(isMaterialized(add))
+
+        let mul = addOrMul(/*useAdd:*/false, a)
+        XCTAssertFalse(isMaterialized(mul))
+        XCTAssertEqual(mul.scalarized(), 25.0);
+        XCTAssertTrue(isMaterialized(mul))
+    }
+
+    func testSimpleLoop() {
+        var sum = Tensor<Float>(0)
+        for i in 1...10 { sum += Float(i) }
+        XCTAssertFalse(isMaterialized(sum))
+        XCTAssertEqual(sum.scalarized(), 55.0, accuracy: 0.00001)
+        XCTAssertTrue(isMaterialized(sum))
+    }
+
+    private func isMaterialized<T: TensorFlowScalar>(_ input: Tensor<T>) -> Bool {
+        let tensor = input.handle.handle
+        guard let lazyTensor = tensor as? LazyTensor else { return true }
+        switch lazyTensor.handle {
+        case .symbolic(let op, _, _): return op.outputs != nil
+        default: return false
+        }
+    }
+
+    static var allTests = [
+        ("testSimpleOperations", testSimpleOperations),
+        ("testMultipleMaterializations", testMultipleMaterializations),
+        ("testSimpleControlFlow", testSimpleControlFlow),
+        ("testSimpleLoop", testSimpleLoop),
+    ]
+}
+

Tests/TensorFlowTests/XCTestManifests.swift

@@ -32,6 +32,7 @@ public func allTests() -> [XCTestCaseEntry] {
         testCase(LazyTensorTraceTests.allTests),
         testCase(LazyTensorOperationTests.allTests),
         testCase(LazyTensorTFFunctionBuilderTests.allTests),
+        testCase(LazyTensorEvaluationTests.allTests),
     ]
 }
 #endif