Adds a mechanism to extract traces explicitly using LazyTensor.

Sources/TensorFlow/Core/LazyTensorTrace.swift

@@ -89,13 +89,55 @@ class LazyTensorTraceBuilder {
         return materializationTraceInfo([lazyOperation])
     }
 
+    /// Returns a trace obtained by tracing the given function.
+    static func trace<In: TensorGroup, Out: TensorGroup>(_ fn: (In) -> Out) -> LazyTensorTrace {
+        precondition(_RuntimeConfig.useLazyTensor, "Lazy tensor is not enabled for tracing.")
+
+        // Set up inputs for running `fn`.
+        let inputOps = In._typeList.map { Self.makePlaceholder(dataType: $0) }
+        let inputHandles = inputOps.map { LazyTensorHandle(_lazy: $0, index: 0) }
+        let input = In(_handles: inputHandles)
+
+        // Run the function.
+        let output: TensorArrayProtocol = fn(input)
+
+        // Set up the closure that determines if a `LazyTensorOperation` should be an output.
+        let outputLazyOperations = output._tensorHandles.map {
+            (handle: _AnyTensorHandle) -> LazyTensorOperation in
+            let lazyOp = lazyTensorOperation(handle)
+            precondition(lazyOp != nil, "Found a non-lazy tensor in output when tracing.")
+            return lazyOp!
+        }
+        let outputIDs = Set<ObjectIdentifier>(
+            outputLazyOperations.lazy.map { ObjectIdentifier($0) })
+
+        // Create the builder and get the trace.
+        let builder = LazyTensorTraceBuilder()
+        builder.neverPromoteConstants = true
+        builder.isOutput = { outputIDs.contains(ObjectIdentifier($0)) }
+        // Set up the inputs for the builder as we need to have them in a specific order.
+        for inputOp in inputOps {
+            builder.updateOperationAndCache(ObjectIdentifier(inputOp), inputOp)
+        }
+        builder.inputs = inputOps
+        for lazyOp in outputLazyOperations { _ = builder.collectLazyOperation(lazyOp) }
+        return LazyTensorTrace(
+            inputs: builder.inputs,
+            operations: builder.operations,
+            outputs: builder.outputs)
+    }
+
     // inputs will be "placeholder" nodes.
     private var inputs: [LazyTensorOperation] = []
     private var inputValues: [TFETensorHandle] = []
     private var operations: [LazyTensorOperation] = []
     private var outputs: [LazyTensorOperation] = []
     private var originalOutputs: [LazyTensorOperation] = []
     private var lazyOpsCache: [ObjectIdentifier: LazyTensorOperation] = [:]
+    /// A flag that controls promotion of constants to inputs.
+    private var neverPromoteConstants: Bool = false
+    /// A closure that determines whether a `LazyTensorOperation` is an output.
+    private var isOutput: (LazyTensorOperation) -> Bool = LazyTensorHandle.isLive
 
     private func updateOperationAndCache(
         _ id: ObjectIdentifier, _ node: LazyTensorOperation
@@ -117,14 +159,25 @@ class LazyTensorTraceBuilder {
         return LazyTensorHandle(_lazy: result, index: 0)
     }
 
-    private func makePlaceholderTensor(
-        with handle: TFETensorHandle
-    ) -> LazyTensorHandle {
-        let cTensorHandle = handle._cTensorHandle
-        let dtype = TensorDataType(TFE_TensorHandleDataType(cTensorHandle))
-        let dtypeAttr = LazyTensorOperation.Attribute.tensorDataTypeValue(dtype)
+    /// Returns the `LazyTensorOperation`, if any, for this handle.
+    private static func lazyTensorOperation(_ handle: _AnyTensorHandle) -> LazyTensorOperation? {
+        guard case let .symbolic(lazyOp, _, _)? = (handle as? LazyTensorHandle)?.handle else {
+            return nil
+        }
+        return lazyOp
+    }
+
+    private static func makePlaceholder(dataType: TensorDataType) -> LazyTensorOperation {
         let placeholder = LazyTensorOperation("Placeholder", 1)
+        let dtypeAttr = LazyTensorOperation.Attribute.tensorDataTypeValue(dataType)
         placeholder.attributes = ["dtype": dtypeAttr]
+        return placeholder
+    }
+
+    private func makePlaceholderTensor(handle: TFETensorHandle) -> LazyTensorHandle {
+        let cTensorHandle = handle._cTensorHandle
+        let dtype = TensorDataType(TFE_TensorHandleDataType(cTensorHandle))
+        let placeholder = Self.makePlaceholder(dataType: dtype)
         updateOperationAndCache(ObjectIdentifier(handle), placeholder)
         inputs.append(placeholder)
         inputValues.append(handle)
@@ -138,12 +191,12 @@ class LazyTensorTraceBuilder {
         if let lazyOp = lazyOpsCache[id] {
             return LazyTensorHandle(_lazy: lazyOp, index: 0)
         }
-        return asConst
+        return asConst || neverPromoteConstants
             ? makeConstTensor(with: handle)
-            : makePlaceholderTensor(with: handle)
+            : makePlaceholderTensor(handle: handle)
     }
 
-    /// Return the original tensor or a concrete tensor that is promoted to a
+    /// Returns the original tensor or a concrete tensor that is promoted to a
     /// placeholder input.
     private func maybePromotedTensor(_ lazyHandle: LazyTensorHandle) -> LazyTensorHandle {
         switch lazyHandle.handle {
@@ -178,13 +231,15 @@ class LazyTensorTraceBuilder {
         if let cachedLazyOp = lazyOpsCache[id] {
             return cachedLazyOp
         }
-
+        precondition(
+            lazyOp.name != "Placeholder",
+            "The operation cannot already be a placeholder.")
         let newLazyOp = LazyTensorOperation(lazyOp.name, lazyOp.outputCount)
         newLazyOp.attributes = lazyOp.attributes
         newLazyOp.inputs = lazyOp.inputs.map { maybePromotedInput($0) }
         updateOperationAndCache(id, newLazyOp)
 
-        if LazyTensorHandle.isLive(lazyOp) {
+        if isOutput(lazyOp) {
             outputs.append(newLazyOp)
             originalOutputs.append(lazyOp)
         }

Tests/TensorFlowTests/LazyTensorExplicitTraceTests.swift

@@ -0,0 +1,152 @@
+// 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 LazyTensorExplicitTraceTests: XCTestCase {
+    override class func setUp() {
+        super.setUp()
+        _RuntimeConfig.useLazyTensor = true
+    }
+
+    override class func tearDown() {
+        super.tearDown()
+        _RuntimeConfig.useLazyTensor = false
+    }
+
+    func testSingleInput() {
+        func fn(x: Tensor<Float>) -> Tensor<Float> { return x + x }
+        let trace = LazyTensorTraceBuilder.trace(fn)
+        XCTAssertEqual(trace.description,
+            """
+            lazyTrace_2(%0: float) -> (%1) {
+              %1 = Add[T: float](%0, %0)
+            }
+            """)
+        let outputs = runTrace(trace: trace, input: Tensor<Float>(10.0))
+        XCTAssertEqual(outputs.count, 1)
+        XCTAssertEqual(outputs[0].valueDescription, "20.0")
+    }
+
+    func testTensorGroupInputOutputs() {
+        typealias TensorFloatInt32Pair = Zip2TensorGroup<Tensor<Float>, Tensor<Int32>>
+        typealias TensorInt32FloatPair = Zip2TensorGroup<Tensor<Int32>, Tensor<Float>>
+        func fn(input: TensorFloatInt32Pair) -> TensorInt32FloatPair {
+            return TensorInt32FloatPair(input.second * 4, input.first + 3.0)
+        }
+        let trace = LazyTensorTraceBuilder.trace(fn)
+        XCTAssertEqual(trace.description,
+            """
+            lazyTrace_6(%0: float, %1: int32) -> (%3, %5) {
+              %2 = Const[dtype: int32, value: 4]()
+              %3 = Mul[T: int32](%1, %2)
+              %4 = Const[dtype: float, value: 3.0]()
+              %5 = Add[T: float](%0, %4)
+            }
+            """)
+        let outputs = runTrace(
+            trace: trace,
+            input: TensorFloatInt32Pair(Tensor<Float>(10.0), Tensor<Int32>(5)))
+        XCTAssertEqual(outputs.count, 2)
+        XCTAssertEqual(outputs[0].valueDescription, "20")
+        XCTAssertEqual(outputs[1].valueDescription, "13.0")
+    }
+
+    func testClosureCapturesOfTensors() {
+        let x = Tensor<Float>(10.0)
+        let y = x + x
+        func fn(input: Tensor<Float>) -> Tensor<Float> {
+            return input * y
+        }
+        let trace = LazyTensorTraceBuilder.trace(fn)
+        /// Note that the computation x + x is encoded in the trace.
+        XCTAssertEqual(trace.description,
+            """
+            lazyTrace_4(%0: float) -> (%3) {
+              %1 = Const[dtype: float, value: 10.0]()
+              %2 = Add[T: float](%1, %1)
+              %3 = Mul[T: float](%0, %2)
+            }
+            """)
+        let outputs = runTrace(
+            trace: trace,
+            input: Tensor<Float>(5.0))
+        XCTAssertEqual(outputs.count, 1)
+        XCTAssertEqual(outputs[0].valueDescription, "100.0")
+    }
+
+    func testClosureCapturesOfNonTensors() {
+        let x: Float = 5.0
+        func fn(input: Tensor<Float>) -> Tensor<Float> {
+            return input * Tensor<Float>(x)
+        }
+        let trace = LazyTensorTraceBuilder.trace(fn)
+        /// Note that the computation x + x is encoded in the trace.
+        XCTAssertEqual(trace.description,
+            """
+            lazyTrace_3(%0: float) -> (%2) {
+              %1 = Const[dtype: float, value: 5.0]()
+              %2 = Mul[T: float](%0, %1)
+            }
+            """)
+        let outputs = runTrace(trace: trace, input: Tensor<Float>(23.0))
+        XCTAssertEqual(outputs.count, 1)
+        XCTAssertEqual(outputs[0].valueDescription, "115.0")
+    }
+
+    func testNestedTracing() {
+        func square(input: Tensor<Float>) -> Tensor<Float> {
+            return input * input
+        }
+
+        func nestedTrace(input: Tensor<Float>) -> Tensor<Float> {
+            let trace = LazyTensorTraceBuilder.trace(square)
+            let outputs = runTrace(trace: trace, input: Tensor<Float>(3.0))
+            XCTAssertEqual(outputs.count, 1)
+            let handle = TensorHandle<Float>(handle: outputs[0])
+            let y = Tensor<Float>(handle: handle)
+            return y + input
+        }
+
+        let trace = LazyTensorTraceBuilder.trace(nestedTrace)
+        XCTAssertEqual(trace.description,
+            """
+            lazyTrace_3(%0: float) -> (%2) {
+              %1 = Const[dtype: float, value: 9.0]()
+              %2 = Add[T: float](%1, %0)
+            }
+            """)
+        let outputs = runTrace(trace: trace, input: Tensor<Float>(4.0))
+        XCTAssertEqual(outputs.count, 1)
+        XCTAssertEqual(outputs[0].valueDescription, "13.0")
+    }
+
+    private func runTrace(trace: LazyTensorTrace, input: TensorGroup) -> [TFETensorHandle] {
+        let tffunc = TFFunction(trace: trace)
+        let inputHandles = input._tensorHandles.map { $0._tfeTensorHandle }
+        let outputHandles = tffunc.execute(inputHandles)
+        return outputHandles
+    }
+
+    static var allTests = [
+        ("testSingleInput", testSingleInput),
+        ("testTensorGroupInputOutputs", testTensorGroupInputOutputs),
+        ("testClosureCapturesOfTensors", testClosureCapturesOfTensors),
+        ("testClosureCapturesOfNonTensors", testClosureCapturesOfNonTensors),
+        ("testNestedTracing", testNestedTracing)
+    ]
+}

Tests/TensorFlowTests/XCTestManifests.swift

@@ -31,6 +31,7 @@ public func allTests() -> [XCTestCaseEntry] {
         testCase(MathOperatorTests.allTests),
         testCase(LazyTensorTests.allTests),
         testCase(LazyTensorTraceTests.allTests),
+        testCase(LazyTensorExplicitTraceTests.allTests),
         testCase(LazyTensorOperationTests.allTests),
         testCase(LazyTensorTFFunctionBuilderTests.allTests),
         testCase(LazyTensorEvaluationTests.allTests),