LazyTensorTFFunctionBuilderTests.swift

// 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 CTensorFlow
import XCTest

@testable import TensorFlow

final class LazyTensorTFFunctionBuilderTests: LazyTensorTestCase {
  func testSingletonInputs() {
    let a = materializedLazyTensor(Tensor<Float>(10.0))
    let w = _Raw.identity(a)
    XCTAssertEqual(
      tfFunction(w, "testSingletonInputs")!.description,
      """

      testSingletonInputs(placeholder_0:float) -> (identity_1:float) {
        Identity_1 = Identity[T=float](placeholder_0)
        return identity_1 = Identity_1:output:0
      }

      """)
  }

  func testListInputs() {
    let a = materializedLazyTensor(Tensor<Float>(10.0))
    let b = materializedLazyTensor(Tensor<Float>(2.0))
    let w = _Raw.addN(inputs: [a, b])
    XCTAssertEqual(
      tfFunction(w, "testListInputs")!.description,
      """

      testListInputs(placeholder_0:float, placeholder_1:float) -> (addn_2:float) {
        AddN_2 = AddN[N=2, T=float](placeholder_0, placeholder_1)
        return addn_2 = AddN_2:sum:0
      }

      """)
  }

  func testSequence() {
    let a = materializedLazyTensor(Tensor<Float>(10.0))
    let b = materializedLazyTensor(Tensor<Float>(2.0))
    let c = materializedLazyTensor(Tensor<Float>(3.0))
    let w = a + b * c
    XCTAssertEqual(
      tfFunction(w, "sequence")!.description,
      """

      sequence(placeholder_0:float, placeholder_1:float, placeholder_2:float) -> (addv2_4:float) {
        Mul_3 = Mul[T=float](placeholder_1, placeholder_2)
        AddV2_4 = AddV2[T=float](placeholder_0, Mul_3:z:0)
        return addv2_4 = AddV2_4:z:0
      }

      """)
  }

  func testAttributes() {
    // If tests ops such as "AttrBool" are available, testing the handing of
    // attributes would be very simple. However, tests ops are not
    // registered into the runtime by default (which is reasonable). If it
    // is possible to get a test-only libtensorflow.so, we should simplify
    // this test usinge the test ops.

    let a = materializedLazyTensor(Tensor<Float>(10.0))
    let b = materializedLazyTensor(Tensor<Float>(20.0))

    // Bool attribute
    let boolAttr = LazyTensorOperation("MatrixInverse", 1)
    boolAttr.updateAttribute("adjoint", true)
    boolAttr.addInput(a)
    XCTAssertEqual(
      tfFunction(boolAttr, "boolAttr").description,
      """

      boolAttr(placeholder_0:float) -> () {
        MatrixInverse_1 = MatrixInverse[T=float, adjoint=true](placeholder_0)
      }

      """)

    // Int attribute
    let intAttr = LazyTensorOperation("Unpack", 1)
    intAttr.updateAttribute("axis", 0)
    intAttr.updateAttribute("num", 1)
    intAttr.updateAttribute("T", Float.tensorFlowDataType)
    intAttr.addInput(a)
    XCTAssertEqual(
      tfFunction(intAttr, "intAttr").description,
      """

      intAttr(placeholder_0:float) -> () {
        Unpack_1 = Unpack[T=float, axis=0, num=1](placeholder_0)
      }

      """)

    // Float attribute
    let floatAttr = LazyTensorOperation("ApproximateEqual", 1)
    floatAttr.updateAttribute("T", Float.tensorFlowDataType)
    floatAttr.updateAttribute("tolerance", 0.01)
    floatAttr.addInput(a)
    floatAttr.addInput(b)
    XCTAssertEqual(
      tfFunction(floatAttr, "floatAttr").description,
      """

      floatAttr(placeholder_0:float, placeholder_1:float) -> () {
        ApproximateEqual_2 = ApproximateEqual[T=float, tolerance=0.01](placeholder_0, placeholder_1)
      }

      """)

    // String attribute
    let stringAttr = LazyTensorOperation("PrintV2", 0)
    let tag = StringTensor("Hello!")
    stringAttr.updateAttribute("output_stream", "stream")
    stringAttr.addInput(tag)
    XCTAssertEqual(
      tfFunction(stringAttr, "stringAttr").description,
      """

      stringAttr() -> () {
        Const_0 = Const[dtype=string, value=Tensor<type: string shape: [] values: Hello!>]()
        PrintV2_1 = PrintV2[end=\"\\n\", output_stream=\"stream\"](Const_0:output:0)
      }

      """)

    // TensorShape attr
    let shapeAttr = LazyTensorOperation("EnsureShape", 1)
    shapeAttr.updateAttribute("shape", TensorShape([5, 6]))
    shapeAttr.updateAttribute("T", Float.tensorFlowDataType)
    shapeAttr.addInput(a)
    XCTAssertEqual(
      tfFunction(shapeAttr, "shapeAttr").description,
      """

      shapeAttr(placeholder_0:float) -> () {
        EnsureShape_1 = EnsureShape[T=float, shape=[5,6]](placeholder_0)
      }

      """)

    // [Int], [TensorShape?] & [TensorDataType] attribute.
    let arrayAttr1 = LazyTensorOperation("PrelinearizeTuple", 0)
    arrayAttr1.updateAttribute(
      "dtypes",
      [Float.tensorFlowDataType, Float.tensorFlowDataType])  // [TensorDataType]
    arrayAttr1.updateAttribute("shapes", [[1, 2], nil])  // [TensorShape?]
    arrayAttr1.updateAttribute("layouts", [3, 4])  // [Int]
    arrayAttr1.addInputList([a, b])

    XCTAssertEqual(
      tfFunction(arrayAttr1, "arrayAttr1").description,
      """

      arrayAttr1(placeholder_0:float, placeholder_1:float) -> () {
        PrelinearizeTuple_2 = PrelinearizeTuple[dtypes={float, float}, layouts=[3, 4], shapes=[[1,2], <unknown>]](placeholder_0, placeholder_1)
      }

      """)

    // Const Tensor attribute.
    let constTensorAttr = LazyTensorOperation("Const", 0)
    let x = Tensor<Float>(5.5)
    constTensorAttr.updateAttribute("dtype", Float.tensorFlowDataType)
    constTensorAttr.updateAttribute("value", x.handle.handle._tfeTensorHandle)
    XCTAssertEqual(
      tfFunction(constTensorAttr, "constTensorAttr").description,
      """

      constTensorAttr() -> () {
        Const_0 = Const[dtype=float, value=Tensor<type: float shape: [] values: 5.5>]()
      }

      """)

    // 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(
    _ lazyOp: LazyTensorOperation,
    _ name: String? = nil
  ) -> TFFunction {
    return TFFunction(trace: lazyTensorTrace(lazyOp), name: name)
  }

  private func materializedLazyTensor<T: TensorFlowScalar>(
    _ input: Tensor<T>
  ) -> Tensor<T> {
    let concreteHandle = input.handle.handle._tfeTensorHandle
    let materializedHandle = LazyTensorHandle(_materialized: concreteHandle)
    return Tensor(handle: TensorHandle<T>(handle: materializedHandle))
  }

  private func tfFunction<T: TensorFlowScalar>(
    _ input: Tensor<T>,
    _ name: String? = nil
  ) -> TFFunction? {
    let tensor = input.handle.handle
    guard let lazyTensor = tensor as? LazyTensorHandle else {
      XCTFail("Trying to get TFFunction for a non-lazy tensor.")
      return nil
    }
    guard case let .symbolic(lazyOp, _, _) = lazyTensor.handle else {
      XCTFail("Cannot get TFFunction for a concrete tensor.")
      return nil
    }
    return TFFunction(trace: lazyTensorTrace(lazyOp), name: name)
  }

  private func lazyTensorTrace(_ lazyOp: LazyTensorOperation) -> LazyTensorTrace {
    let traceInfo = LazyTensorTraceBuilder.materializationTraceInfo(lazyOp)
    return traceInfo.trace
  }

  static var allTests = [
    ("testSingletonInputs", testSingletonInputs),
    ("testListInputs", testListInputs),
    ("testSequence", testSequence),
    ("testAttributes", testAttributes),
  ]
}