Merge branch 'tensorflow' into tensorflow-merge

Author: dan-zheng

include/swift/AST/AutoDiff.h

@@ -529,12 +529,18 @@ template <typename T> struct DenseMapInfo;
 template <> struct DenseMapInfo<AutoDiffConfig> {
   static AutoDiffConfig getEmptyKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getEmptyKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
             nullptr};
   }
 
   static AutoDiffConfig getTombstoneKey() {
     auto *ptr = llvm::DenseMapInfo<void *>::getTombstoneKey();
+    // The `derivativeGenericSignature` component must be `nullptr` so that
+    // `getHashValue` and `isEqual` do not try to `getCanonicalSignature()` on
+    // an invalid pointer.
     return {static_cast<IndexSubset *>(ptr), static_cast<IndexSubset *>(ptr),
             nullptr};
   }

test/AutoDiff/derivative_attr_type_checking.swift

@@ -532,3 +532,33 @@ struct StoredProperty: Differentiable {
     (stored, { _ in .zero })
   }
 }
+
+// When the generic signature was not considered while calculating the actual pullback type, the
+// typechecker did not realize that `T.TangentVector == Float`, and therefore it complained that
+// "'pullback' does not have expected type '(Float) -> (Float)'".
+// Users were able to work around this by setting the pullback type to `(Float) -> Float`.
+func genericSignatureConsidered<T>(_ x: T) -> T { fatalError() }
+@derivative(of: genericSignatureConsidered)
+func dGenericSignatureConsidered<T>(_ x: T)
+  -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
+  where T: Differentiable, T.TangentVector == Float
+{
+  fatalError()
+}
+
+// When the generic signature was not considered while calculating the actual pullback type,
+// the typechecker complained that the pullback type was not correct, and there was no pullback type
+// that users could specify to satisfy the typechecker.
+struct Wrapper<T: AdditiveArithmetic & Equatable>: AdditiveArithmetic, Equatable {
+  var t: T
+  init(_ t: T) { self.t = t }
+}
+extension Wrapper: Differentiable where T: Differentiable, T == T.TangentVector {
+  typealias TangentVector = Wrapper<T.TangentVector>
+}
+extension Wrapper where T: Differentiable, T == T.TangentVector {
+  @derivative(of: init(_:))
+  static func dInit(_ t: T) -> (value: Self, pullback: (Wrapper<T>.TangentVector) -> (T)) {
+    fatalError()
+  }
+}

test/AutoDiff/derivative_registration.swift

@@ -156,4 +156,24 @@ DerivativeRegistrationTests.testWithLeakChecking("DerivativeGenericSignature") {
   expectEqual(1000, dx)
 }
 
+// When non-canonicalized generic signatures are used to compare derivative configurations, the
+// `@differentiable` and `@derivative` attributes create separate derivatives, and we get a
+// duplicate symbol error in TBDGen.
+public protocol RefinesDifferentiable: Differentiable {}
+extension Float: RefinesDifferentiable {}
+@differentiable(where T: Differentiable, T: RefinesDifferentiable)
+public func nonCanonicalizedGenSigComparison<T>(_ t: T) -> T { t }
+@derivative(of: nonCanonicalizedGenSigComparison)
+public func dNonCanonicalizedGenSigComparison<T: RefinesDifferentiable>(_ t: T)
+  -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
+{
+  (t, { _ in T.TangentVector.zero })
+}
+DerivativeRegistrationTests.testWithLeakChecking("NonCanonicalizedGenericSignatureComparison") {
+  let dx = gradient(at: Float(0), in: nonCanonicalizedGenSigComparison)
+  // Expect that we use the custom registered derivative, not a generated derivative (which would
+  // give a gradient of 1).
+  expectEqual(0, dx)
+}
+
 runAllTests()

utils/update_checkout/update-checkout-config.json

@@ -296,7 +296,7 @@
                 "indexstore-db": "swift-DEVELOPMENT-SNAPSHOT-2019-12-16-a",
                 "sourcekit-lsp": "swift-DEVELOPMENT-SNAPSHOT-2019-12-16-a",
                 "tensorflow": "7c7d924821a8b1b20433c2f3f484bbd409873a84",
-                "tensorflow-swift-apis": "47de23315b574f936daefe1e781bbfd5dfe1a247",
+                "tensorflow-swift-apis": "0aaf944b91703ea6ae9062b3a720d3e27b7eba36",
                 "tensorflow-swift-quote": "34d4112c294e9eccfb575dd8f39e6e3b0b6bc008"
             }
         }