Add support for tensorFunctionPointer attribute in LazyTensor.

Sources/TensorFlow/Core/LazyTensorOperation.swift

@@ -388,8 +388,7 @@ extension LazyTensorOperation: TFTensorOperation {
 
     func updateAttribute<In: TensorGroup, Out: TensorGroup>(
         _ name: String, _ value: (In) -> Out) {
-        // TODO:
-        fatalError("Unimplemented [TFFunction] attribute.")
+        updateAttribute(name, _TensorFunctionPointer(name: _tffunc(value)))
     }
 
     func execute() {
@@ -443,7 +442,7 @@ extension LazyTensorOperation: TFTensorOperation {
             case .tensorDataTypeArray(let v): op.updateAttribute(name, v)
             case .optionalTensorShape(let v): op.updateAttribute(name, v)
             case .optionalTensorShapeArray(let v): op.updateAttribute(name, v)
-            case .tensorFunctionPointer(_): fatalError("tensorFunctionPointer Unimplemented!")
+            case .tensorFunctionPointer(let v): op.updateAttribute(name, v)
             }
         }
         op.execute()

Sources/TensorFlow/Core/LazyTensorShapeInference.swift

@@ -37,7 +37,7 @@ extension LazyTensorOperation {
             case .tensorDataTypeArray(let v): op.updateAttribute(name, v)
             case .optionalTensorShape(let v): op.updateAttribute(name, v)
             case .optionalTensorShapeArray(let v): op.updateAttribute(name, v)
-            case .tensorFunctionPointer(_): fatalError("tensorFunctionPointer Unimplemented!")
+            case .tensorFunctionPointer(let v): op.updateAttribute(name, v)
             }
         }
         return op

Sources/TensorFlow/Core/LazyTensorTFFunctionBuilder.swift

@@ -163,6 +163,8 @@ class TFGraph {
                     }
                 }
             }
+        case .tensorFunctionPointer(let value):
+            TF_SetAttrFuncName(description, name, value.name, value.name.count)
         default: fatalError("Unhandled attribute \(name):\(attribute)")
         }
     }

Tests/TensorFlowTests/LazyTensorTFFunctionBuilderTests.swift

@@ -198,6 +198,22 @@ final class LazyTensorTFFunctionBuilderTests : XCTestCase {
 
             """)
 
+        // TensorFunctionPointer attribute.
+        let statelessWhile = LazyTensorOperation("StatelessWhile", 1)
+        statelessWhile.updateAttribute("T", [Float.tensorFlowDataType])
+        statelessWhile.updateAttribute("cond", _TensorFunctionPointer(name: "cond"))
+        statelessWhile.updateAttribute("body", _TensorFunctionPointer(name: "body"))
+        statelessWhile.addInputList([a])
+        XCTAssertEqual(
+            tfFunction(statelessWhile, "statelessWhile").description,
+            """
+
+            statelessWhile(placeholder_0:float) -> () {
+              StatelessWhile_1 = StatelessWhile[T={float}, body=body, cond=cond, output_shapes=[], parallel_iterations=10](placeholder_0)
+            }
+
+            """)
+
     }
 
     private func tfFunction(

Tests/TensorFlowTests/LazyTensorTraceTests.swift

@@ -206,6 +206,36 @@ final class LazyTensorTraceTests: XCTestCase {
         XCTAssertEqual(z.scalarized(), 9.0)
     }
 
+    func testTraceWithFunctionAttributes() {
+        typealias Int32Pair = Zip2TensorGroup<Tensor<Int32>, Tensor<Int32>>
+        func thenBranch(x: Tensor<Float>) -> Tensor<Float> {
+            return x + 10.0
+        }
+        func elseBranch(x: Tensor<Float>) -> Tensor<Float> {
+            return x - 9.0
+        }
+        let c: Tensor<Float> = Raw.if_(
+            cond: Tensor<Bool>(false),
+            Tensor<Float>(20.0),
+            thenBranch: thenBranch,
+            elseBranch: elseBranch,
+            outputShapes: [nil])
+        let cLazyOp = lazyTensorOperation(c)!
+        let cTraceInfo = LazyTensorTraceBuilder.materializationTraceInfo(cLazyOp)
+        let cTrace = cTraceInfo.trace
+        XCTAssertEqual(
+            cTrace.description,
+            """
+            lazyTrace_3() -> (%2) {
+              %0 = Const[dtype: bool, value: false]()
+              %1 = Const[dtype: float, value: 20.0]()
+              %2 = If[Tcond: bool, Tin: [float], Tout: [float], else_branch: TFFunction(lazyTrace_3_kMDsaAFRUp8), output_shapes: [nil], then_branch: TFFunction(lazyTrace_3_sayLTaDTeLE)](%0, [%1])
+            }
+            """)
+        // Returns the result of the else branch.
+        XCTAssertEqual(c.scalarized(), 11.0)
+    }
+
     private func lazyTensorOperation<T: TensorFlowScalar>(
         _ input: Tensor<T>
     ) -> LazyTensorOperation? {
@@ -236,6 +266,7 @@ final class LazyTensorTraceTests: XCTestCase {
         ("testMultipleTargets", testMultipleTargets),
         ("testSimpleControlFlow", testSimpleControlFlow),
         ("testManualConstPromotion", testManualConstPromotion),
-        ("testConstPromotion", testConstPromotion)
+        ("testConstPromotion", testConstPromotion),
+        ("testTraceWithFunctionAttributes", testTraceWithFunctionAttributes)
     ]
 }