From 486e0548f76bbc975a449ea6fea6f225d1357116 Mon Sep 17 00:00:00 2001
From: Brett Koonce <koonce@hello.com>
Date: Wed, 19 Jun 2019 16:06:30 -0500
Subject: [PATCH] transformer: upstream api changes

---
 Transformer/Model.swift     | 10 +++++-----
 Transformer/Operators.swift |  4 ++--
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/Transformer/Model.swift b/Transformer/Model.swift
index 3fcf71ccd50..70633007606 100644
--- a/Transformer/Model.swift
+++ b/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)
diff --git a/Transformer/Operators.swift b/Transformer/Operators.swift
index ed16aa215ac..20944d45bd9 100644
--- a/Transformer/Operators.swift
+++ b/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,