Handle input_tensors during shape inference. (#453)

Author: bgogul

Diff for: Sources/TensorFlow/Core/LazyTensorShapeInference.swift

@@ -66,6 +66,36 @@ extension LazyTensorOperation {
             }
         }
 
+        // Returns the `CTensor`, selectively materializing it if needed.
+        func cTensor(handle: LazyTensorHandle) -> CTensor? {
+            switch handle.handle {
+            case .concrete(let h, _):
+                let cTensor = TFE_TensorHandleResolve(h._cTensorHandle, status)
+                checkOk(status)
+                return cTensor
+            case .symbolic(let op, _, _):
+                // TODO(https://bugs.swift.org/browse/TF-765): "Pack" is used
+                // for creating tensors from array literals. So, allow
+                // materialization for 'Pack' so that we can get the shape for
+                // array literals. We should revisit this heuristic.
+                if op.name != "Pack" { return nil }
+                let cTensor = TFE_TensorHandleResolve(handle._cTensorHandle, status)
+                checkOk(status)
+                return cTensor
+            }
+        }
+
+        // Create `inputTensors` consisting of *only* materialized inputs.
+        var inputTensors: [CTensor?] = []
+        for input in inputs {
+            switch input {
+            case .single(let v):
+                inputTensors.append(cTensor(handle: v))
+            case .list(let values):
+                inputTensors.append(contentsOf: values.lazy.map { cTensor(handle: $0) } )
+            }
+        }
+
         // This will be filled in by `TFE_InferShapes` and should be freed later.
         var outputShapeListPtr = UnsafeMutablePointer<TF_ShapeAndTypeList>(nil)
         defer { TF_DeleteShapeAndTypeList(outputShapeListPtr) }
@@ -76,17 +106,18 @@ extension LazyTensorOperation {
             TF_DeleteStatus(tfeOp.status)
         }
 
-        TFE_InferShapes(
-            tfeOp.op,
-            /*input_shapes*/ inputShapeList,
-            /*input_tensors*/ nil,
-            /*input_tensors_as_shapes*/ nil,
-            /*input_resource_shapes_and_types*/ nil,
-            /*output_shapes*/ &outputShapeListPtr,
-            /*output_resource_shapes_and_types*/ nil,
-            status)
-
-        checkOk(status)
+        inputTensors.withUnsafeMutableBufferPointer { buffer in
+            TFE_InferShapes(
+                tfeOp.op,
+                /*input_shapes*/ inputShapeList,
+                /*input_tensors*/ buffer.baseAddress!,
+                /*input_tensors_as_shapes*/ nil,
+                /*input_resource_shapes_and_types*/ nil,
+                /*output_shapes*/ &outputShapeListPtr,
+                /*output_resource_shapes_and_types*/ nil,
+                status)
+            checkOk(status)
+        }
 
         precondition(outputShapeListPtr != nil, "TFE_InferShapes returned nil for output shapes")
         let outputShapeList = outputShapeListPtr!.pointee

Diff for: Tests/TensorFlowTests/LazyTensorShapeInferenceTests.swift

@@ -64,8 +64,38 @@ final class LazyTensorShapeInferenceTests: XCTestCase {
         XCTAssertTrue(xLazyTensorOperation.isMaterialized)
     }
 
+    /// Checks scenarios where shapes are computed from input tensors.
+    func testShapeComputationsWithInputTensors() {
+        let a = Tensor<Float>(shape: [3, 1], scalars: [1.0, 2.0, 3.0])
+        let b = a.reshaped(toShape: [1, 3])
+
+        let bLazyTensorOperation = b._lazyTensor!.lazyTensorOperation!
+        XCTAssertFalse(bLazyTensorOperation.isMaterialized)
+
+        let bShape = b.shape
+        XCTAssertEqual(bShape.rank, 2)
+        XCTAssertEqual(bShape.dimensions, [1, 3])
+        XCTAssertFalse(bLazyTensorOperation.isMaterialized)
+
+        let c = Tensor<Float>(repeating: 5, shape: [4, 5, 6])
+        let cLazyTensorOperation = c._lazyTensor!.lazyTensorOperation!
+        XCTAssertFalse(cLazyTensorOperation.isMaterialized)
+
+        let cShape = c.shape
+        XCTAssertEqual(cShape.rank, 3)
+        XCTAssertEqual(cShape.dimensions, [4, 5, 6])
+        XCTAssertFalse(cLazyTensorOperation.isMaterialized)
+
+        // Trigger materialization.
+        let _ = b._rawTensorHandle
+        let _ = c._rawTensorHandle
+        XCTAssertTrue(bLazyTensorOperation.isMaterialized)
+        XCTAssertTrue(cLazyTensorOperation.isMaterialized)
+    }
+
     static var allTests = [
-        ("testSimpleShapeComputations", testSimpleShapeComputations)
+        ("testSimpleShapeComputations", testSimpleShapeComputations),
+        ("testShapeComputationsWithInputTensors", testShapeComputationsWithInputTensors)
     ]
 }