transformer: upstream api changes (#171)

Author: brettkoonce

Transformer/Model.swift

@@ -51,9 +51,9 @@ struct FeedForward: Layer {
 }
 
 struct AttentionInput: Differentiable {
-    let query: Tensor<Float>
-    let key: Tensor<Float>
-    let value: Tensor<Float>
+    var query: Tensor<Float>
+    var key: Tensor<Float>
+    var value: Tensor<Float>
 }
 
 @differentiable(wrt: (query, key, value), vjp: _vjpMakeAttentionInput)
@@ -69,8 +69,8 @@ func _vjpMakeAttentionInput(query: Tensor<Float>, key: Tensor<Float>, value: Ten
 }
 
 struct AttentionContext: Differentiable {
-    let key: Tensor<Float>
-    let value: Tensor<Float>
+    var key: Tensor<Float>
+    var value: Tensor<Float>
 }
 
 @differentiable(wrt: (key, value), vjp: _vjpMakeAttentionContext)

Transformer/Operators.swift

@@ -29,7 +29,7 @@ func gelu<Scalar: TensorFlowFloatingPoint>(_ x: Tensor<Scalar>) -> Tensor<Scalar
 @differentiable(
     wrt: (left, right),
     vjp: _vjpBatchedMatmul
-    where Scalar : Differentiable & FloatingPoint
+    where Scalar : Differentiable & TensorFlowFloatingPoint
 )
 func batchedMatmul<Scalar : Numeric>(
     _ left: Tensor<Scalar>,
@@ -41,7 +41,7 @@ func batchedMatmul<Scalar : Numeric>(
 }
 
 @usableFromInline
-func _vjpBatchedMatmul<Scalar : Differentiable & FloatingPoint>(
+func _vjpBatchedMatmul<Scalar : Differentiable & TensorFlowFloatingPoint>(
     _ left: Tensor<Scalar>,
     _ right: Tensor<Scalar>,
     adjointLeft: Bool,