tensorflow/swift-apis compilation is painfully slow. It's probably due to type-checking.

Swift compilation debugging tips here. swift build -Xswiftc -Xfrontend -Xswiftc -debug-time-function-bodies results below.

# Worst offenders, time in milliseconds.
(1767, "global function 'gelu'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1225:13')
(1767, "global function 'gelu'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1225:13')
(1767, "global function 'gelu'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1225:13')
(1866, "instance method 'sha512()'", '/Users/danielzheng/swift-apis/Sources/Tensor/TensorUtilities.swift:149:10')
(2220, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:152:17')
(2220, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:152:17')
(2220, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:152:17')
(2304, "global function 'root'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1375:13')
(2304, "global function 'root'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1375:13')
(2304, "global function 'root'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1375:13')
(2528, "global function 'hingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:109:13')
(2528, "global function 'hingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:109:13')
(2528, "global function 'hingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:109:13')
(3796, "global function 'categoricalHingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:139:13')
(3796, "global function 'categoricalHingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:139:13')
(3796, "global function 'categoricalHingeLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:139:13')
(4102, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:429:17')
(4102, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:429:17')
(4102, "instance method 'update(_:along:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Optimizers/MomentumBased.swift:429:17')
(4128, "global function 'cosineSimilarity'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1492:13')
(4128, "global function 'cosineSimilarity'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1492:13')
(4128, "global function 'cosineSimilarity'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Operators/Math.swift:1492:13')
(15152, "global function 'logCoshLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:157:13')
(15152, "global function 'logCoshLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:157:13')
(15152, "global function 'logCoshLoss(predicted:expected:reduction:)'", '/Users/danielzheng/swift-apis/Sources/TensorFlow/Loss.swift:157:13')

public func logCoshLoss<Scalar: TensorFlowFloatingPoint>(
    predicted: Tensor<Scalar>,
    expected: Tensor<Scalar>,
    reduction: @differentiable (Tensor<Scalar>) -> Tensor<Scalar> = _mean
) -> Tensor<Scalar> {
    let x = predicted - expected
    return reduction(x + softplus(Tensor(-2) * x) - log(Tensor(2)))
}

swift build -Xswiftc -Xfrontend -Xswiftc -debug-time-compilation didn't print any output for me for some reason.

---

Idea from @allevato: provide a contextual type for integer/float literals to help the type-checker.

It works:

import TensorFlow

@differentiable
public func logCoshLoss<Scalar: TensorFlowFloatingPoint>(
    predicted: Tensor<Scalar>,
    expected: Tensor<Scalar>,
    reduction: @differentiable (Tensor<Scalar>) -> Tensor<Scalar> = _mean
) -> Tensor<Scalar> {
    let x = predicted - expected
    // Original code.
    return reduction(x + softplus(Tensor(-2) * x) - log(Tensor(2)))
}

@differentiable
public func logCoshLossTest<Scalar: TensorFlowFloatingPoint>(
    predicted: Tensor<Scalar>,
    expected: Tensor<Scalar>,
    reduction: @differentiable (Tensor<Scalar>) -> Tensor<Scalar> = _mean
) -> Tensor<Scalar> {
    let x = predicted - expected
    // Tony's suggestion: provide contextual type for literals.
    return reduction(x + softplus(Tensor(-2 as Scalar) * x) - log(Tensor(2 as Scalar)))
}

$ swift -Xfrontend -debug-time-function-bodies timing.swift
timing.swift:4:13: warning: global function 'logCoshLoss(predicted:expected:reduction:)' took 7400ms to type-check (limit: 1ms)
public func logCoshLoss<Scalar: TensorFlowFloatingPoint>(
            ^
timing.swift:15:13: warning: global function 'logCoshLossTest(predicted:expected:reduction:)' took 134ms to type-check (limit: 1ms)
public func logCoshLossTest<Scalar: TensorFlowFloatingPoint>(
            ^

I believe this should not be necessary, and may be a deficiency in the Swift type checker, specifically bidirectional type-checking and constraint propagation. When a contextual type exists, constraints should propagate from outer->inner: so Scalar is the only possible type for the literals 2 and -2 in logCoshLoss. Inner->outer propagation results in a huge disjunction constraint and huge slowdown.

I'm not sure why exactly logCoshLoss takes so long to type-check. I plan to write a minimal reproducer and ask on the forums.

---

Suggestion from @rxwei: splitting tensorflow/swift-apis into more files can help multithreaded compilation, since one thread is spawned per file. Some files like Math.swift are huge and can be split.