生の TensorFlow 演算子

TensorFlow 上に構築された Swift for TensorFlow は、API 設計に新しいアプローチを採用しています。 API は確立されたライブラリから慎重に厳選され、新しい言語イディオムと組み合わせられます。これは、すべての TensorFlow API が Swift API として直接利用できるわけではないことを意味し、API のキュレーションを進化させるには時間と献身的な努力が必要です。ただし、お気に入りの TensorFlow オペレーターが Swift で利用できない場合でも心配する必要はありません。TensorFlow Swift ライブラリを使用すると、 _Raw名前空間でほとんどの TensorFlow オペレーターに透過的にアクセスできます。

TensorFlowをインポートして開始します。

import TensorFlow

生の演算子の呼び出し

コード補完を使用して、 _Raw名前空間で必要な関数を見つけるだけです。

print(_Raw.mul(Tensor([2.0, 3.0]), Tensor([5.0, 6.0])))

[10.0, 18.0]

新しい乗算演算子の定義

Multiply はTensor上で演算子*としてすでに利用可能ですが、それを.*という新しい名前で利用できるようにしたいと考えてみましょう。 Swift では、 extension宣言を使用して既存の型にメソッドや計算されたプロパティを遡及的に追加できます。

次に、拡張子を宣言してTensorに.*を追加し、Tensor のScalar型がNumericに準拠する場合に利用できるようにしましょう。

infix operator .* : MultiplicationPrecedence

extension Tensor where Scalar: Numeric {
    static func .* (_ lhs: Tensor, _ rhs: Tensor) -> Tensor {
        return _Raw.mul(lhs, rhs)
    }
}

let x: Tensor<Double> = [[1.0, 2.0], [3.0, 4.0]]
let y: Tensor<Double> = [[8.0, 7.0], [6.0, 5.0]]
print(x .* y)

[[ 8.0, 14.0],
 [18.0, 20.0]]

ラップされた関数の導関数の定義

生の TensorFlow オペレーター用の Swift API を簡単に定義できるだけでなく、Swift のファーストクラスの自動微分で動作するように微分可能にすることもできます。

.*を微分可能にするには、導関数で@derivative属性を使用し、元の関数をof:ラベルの下の属性引数として指定します。 .*演算子はジェネリック型Scalar Numericに準拠するときに定義されるため、 Tensor<Scalar>をDifferentiableプロトコルに準拠させるには十分ではありません。型安全性を備えて生まれた Swift は、 Scalar TensorFlowFloatingPointプロトコルに準拠することを要求するための@differentiable属性に汎用制約を追加することを思い出させます。これにより、 Tensor<Scalar> Differentiableに準拠します。

@differentiable(where Scalar: TensorFlowFloatingPoint)

infix operator .* : MultiplicationPrecedence

extension Tensor where Scalar: Numeric {
    @differentiable(where Scalar: TensorFlowFloatingPoint)
    static func .* (_ lhs: Tensor,  _ rhs: Tensor) -> Tensor {
        return _Raw.mul(lhs, rhs)
    }
}

extension Tensor where Scalar : TensorFlowFloatingPoint { 
    @derivative(of: .*)
    static func multiplyDerivative(
        _ lhs: Tensor, _ rhs: Tensor
    ) -> (value: Tensor, pullback: (Tensor) -> (Tensor, Tensor)) {
        return (lhs * rhs, { v in
            ((rhs * v).unbroadcasted(to: lhs.shape),
            (lhs * v).unbroadcasted(to: rhs.shape))
        })
    }
}

// Now, we can take the derivative of a function that calls `.*` that we just defined.
print(gradient(at: x, y) { x, y in
    (x .* y).sum()
})

(0.0, 0.0)

他の例

let matrix = Tensor<Float>([[1, 2], [3, 4]])

print(_Raw.matMul(matrix, matrix, transposeA: true, transposeB: true))
print(_Raw.matMul(matrix, matrix, transposeA: true, transposeB: false))
print(_Raw.matMul(matrix, matrix, transposeA: false, transposeB: true))
print(_Raw.matMul(matrix, matrix, transposeA: false, transposeB: false))

[[ 7.0, 15.0],
 [10.0, 22.0]]
[[10.0, 14.0],
 [14.0, 20.0]]
[[ 5.0, 11.0],
 [11.0, 25.0]]
[[ 7.0, 10.0],
 [15.0, 22.0]]