---

yaml
---
r: 338939
b: refs/heads/rxwei-patch-1
c: 7c6d2fa
h: refs/heads/master
i:
  338937: e846ad6
  338935: b6e2678

Author: dan-zheng

[refs]

@@ -1015,7 +1015,7 @@ refs/tags/swift-DEVELOPMENT-SNAPSHOT-2018-08-18-a: b10b1fce14385faa6d44f6b933e95
 refs/heads/rdar-43033749-fix-batch-mode-no-diags-swift-5.0-branch: a14e64eaad30de89f0f5f0b2a782eed7ecdcb255
 refs/heads/revert-19006-error-bridging-integer-type: 8a9065a3696535305ea53fe9b71f91cbe6702019
 refs/heads/revert-19050-revert-19006-error-bridging-integer-type: ecf752d54b05dd0a20f510f0bfa54a3fec3bcaca
-refs/heads/rxwei-patch-1: b14f205d0a0907c06d6ae5d0501ede2f6e3756c8
+refs/heads/rxwei-patch-1: 7c6d2fa0b744f651fe88424628bd4363d54a688a
 refs/heads/shahmishal-patch-1: e58ec0f7488258d42bef51bc3e6d7b3dc74d7b2a
 refs/heads/typelist-existential: 4046359efd541fb5c72d69a92eefc0a784df8f5e
 refs/tags/swift-4.2-DEVELOPMENT-SNAPSHOT-2018-08-20-a: 4319ba09e4fb8650ee86061075c74a016b6baab9

branches/rxwei-patch-1/lib/Sema/TypeCheckAttr.cpp

@@ -3053,15 +3053,6 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   // Set the checked differentiation parameter indices in the attribute.
   attr->setParameterIndices(checkedWrtParamIndices);
 
-  auto insertion =
-      ctx.DifferentiableAttrs.try_emplace({D, checkedWrtParamIndices}, attr);
-  // `@differentiable` attributes are uniqued by their parameter indices.
-  // Reject duplicate attributes for the same decl and parameter indices pair.
-  if (!insertion.second && insertion.first->getSecond() != attr) {
-    diagnoseAndRemoveAttr(attr, diag::differentiable_attr_duplicate);
-    return;
-  }
-
   // Check that original function's result type conforms to `Differentiable`.
   if (whereClauseGenEnv) {
     auto originalResultInterfaceType = !originalResultTy->hasTypeParameter()
@@ -3079,9 +3070,7 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
   }
 
   // Checks that the `candidate` function type equals the `required` function
-  // type, disregarding parameter labels.
-  //
-  // Precondition: `required` has no parameter labels.
+  // type, disregarding parameter labels and tuple result labels.
   std::function<bool(CanAnyFunctionType, CanType)> checkFunctionSignature;
   checkFunctionSignature = [&](CanAnyFunctionType required,
                                CanType candidate) -> bool {
@@ -3096,21 +3085,31 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
         required.getOptGenericSignature())
       return false;
 
-    // Check that parameter types match (disregards labels).
-    if (candidateFnTy.getParams().size() != required.getParams().size())
+    // Check that parameter types match, disregarding labels.
+    if (!std::equal(required.getParams().begin(), required.getParams().end(),
+                    candidateFnTy.getParams().begin(),
+                    [](AnyFunctionType::Param x, AnyFunctionType::Param y) {
+                      return x.getPlainType()->isEqual(y.getPlainType());
+                    }))
       return false;
-    for (auto paramPair : llvm::zip(candidateFnTy.getParams(),
-                                    required.getParams()))
-      if (!std::get<0>(paramPair).getPlainType()->isEqual(
-          std::get<1>(paramPair).getPlainType()))
-        return false;
 
-    // If required result type is non-function, check that result types match
-    // exactly.
+    // If required result type is non-function, check that result types match.
+    // If result types are tuple types, ignore labels.
     CanAnyFunctionType requiredResultFnTy =
         dyn_cast<AnyFunctionType>(required.getResult());
-    if (!requiredResultFnTy)
-      return required.getResult() == candidateFnTy.getResult();
+    if (!requiredResultFnTy) {
+      auto requiredResultTupleTy = required.getResult()->getAs<TupleType>();
+      auto candidateResultTupleTy =
+          candidateFnTy.getResult()->getAs<TupleType>();
+      if (!requiredResultTupleTy || !candidateResultTupleTy)
+        return required.getResult()->isEqual(candidateFnTy.getResult());
+      // If result types are tuple types, check that element types match,
+      // ignoring labels.
+      return std::equal(requiredResultTupleTy->getElementTypes().begin(),
+                        requiredResultTupleTy->getElementTypes().end(),
+                        candidateResultTupleTy->getElementTypes().begin(),
+                        [](Type x, Type y) { return x->isEqual(y); });
+    }
 
     // Required result type is a function. Recurse.
     return checkFunctionSignature(requiredResultFnTy,
@@ -3168,6 +3167,15 @@ void AttributeChecker::visitDifferentiableAttr(DifferentiableAttr *attr) {
     // Memorize the vjp reference in the attribute.
     attr->setVJPFunction(vjp);
   }
+
+  auto insertion =
+      ctx.DifferentiableAttrs.try_emplace({D, checkedWrtParamIndices}, attr);
+  // `@differentiable` attributes are uniqued by their parameter indices.
+  // Reject duplicate attributes for the same decl and parameter indices pair.
+  if (!insertion.second && insertion.first->getSecond() != attr) {
+    diagnoseAndRemoveAttr(attr, diag::differentiable_attr_duplicate);
+    return;
+  }
 }
 
 // SWIFT_ENABLE_TENSORFLOW
@@ -3489,33 +3497,50 @@ void AttributeChecker::visitDifferentiatingAttr(DifferentiatingAttr *attr) {
         derivativeRequirements.push_back(req);
   }
 
-  // Add the derivative to a `@differentiable` attribute on the original
-  // function with the same differentiation parameters. If no such
-  // `@differentiable` attribute exists, create one.
+  // Try to find a `@differentiable` attribute on the original function with the
+  // same differentiation parameters.
   DifferentiableAttr *da = nullptr;
   for (auto *cda : originalFn->getAttrs().getAttributes<DifferentiableAttr>())
     if (checkedWrtParamIndices == cda->getParameterIndices())
       da = const_cast<DifferentiableAttr *>(cda);
+  // If the original function does not have a `@differentiable` attribute with
+  // the same differentiation parameters, create one.
   if (!da) {
     da = DifferentiableAttr::create(ctx, /*implicit*/ true, attr->AtLoc,
                                     attr->getRange(), checkedWrtParamIndices,
-                                    None, None, derivativeRequirements);
+                                    /*jvp*/ None, /*vjp*/ None,
+                                    derivativeRequirements);
+    switch (kind) {
+    case AutoDiffAssociatedFunctionKind::JVP:
+      da->setJVPFunction(derivative);
+      break;
+    case AutoDiffAssociatedFunctionKind::VJP:
+      da->setVJPFunction(derivative);
+      break;
+    }
     auto insertion = ctx.DifferentiableAttrs.try_emplace(
         {originalFn, checkedWrtParamIndices}, da);
-    // `@differentiable` attributes are uniqued by their parameter indices.
-    // Reject duplicate attributes for the same decl and parameter indices pair.
+    // Valid `@differentiable` attributes are uniqued by their parameter
+    // indices. Reject duplicate attributes for the same decl and parameter
+    // indices pair.
     if (!insertion.second && insertion.first->getSecond() != da) {
       diagnoseAndRemoveAttr(da, diag::differentiable_attr_duplicate);
       return;
     }
     originalFn->getAttrs().add(da);
+    return;
   }
-  // Check if the `@differentiable` attribute already has a registered
-  // derivative. If so, emit an error on the `@differentiating` attribute.
-  // Otherwise, register the derivative in the `@differentiable` attribute.
+  // If the original function has a `@differentiable` attribute with the same
+  // differentiation parameters, check if the `@differentiable` attribute
+  // already has a different registered derivative. If so, emit an error on the
+  // `@differentiating` attribute. Otherwise, register the derivative in the
+  // `@differentiable` attribute.
   switch (kind) {
   case AutoDiffAssociatedFunctionKind::JVP:
-    if (da->getJVP() || da->getJVPFunction()) {
+    // If there's a different registered derivative, emit an error.
+    if ((da->getJVP() &&
+         da->getJVP()->Name.getBaseName() != derivative->getBaseName()) ||
+        (da->getJVPFunction() && da->getJVPFunction() != derivative)) {
       diagnoseAndRemoveAttr(
           attr, diag::differentiating_attr_original_already_has_derivative,
           originalFn->getFullName());
@@ -3524,7 +3549,10 @@ void AttributeChecker::visitDifferentiatingAttr(DifferentiatingAttr *attr) {
     da->setJVPFunction(derivative);
     break;
   case AutoDiffAssociatedFunctionKind::VJP:
-    if (da->getVJP() || da->getVJPFunction()) {
+    // If there's a different registered derivative, emit an error.
+    if ((da->getVJP() &&
+         da->getVJP()->Name.getBaseName() != derivative->getBaseName()) ||
+        (da->getVJPFunction() && da->getVJPFunction() != derivative)) {
       diagnoseAndRemoveAttr(
           attr, diag::differentiating_attr_original_already_has_derivative,
           originalFn->getFullName());

branches/rxwei-patch-1/test/AutoDiff/differentiable_attr_type_checking.swift

@@ -505,6 +505,41 @@ func vjpWhere1<T : Differentiable>(x: T) -> (T, (T.CotangentVector) -> T.Cotange
   return (x, { v in v })
 }
 
+// Test derivative functions with result tuple type labels.
+@differentiable(jvp: jvpResultLabels, vjp: vjpResultLabels)
+func derivativeResultLabels(_ x: Float) -> Float {
+  return x
+}
+func jvpResultLabels(_ x: Float) -> (value: Float, differential: (Float) -> Float) {
+  return (x, { $0 })
+}
+func vjpResultLabels(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  return (x, { $0 })
+}
+struct ResultLabelTest {
+  @differentiable(jvp: jvpResultLabels, vjp: vjpResultLabels)
+  static func derivativeResultLabels(_ x: Float) -> Float {
+    return x
+  }
+  static func jvpResultLabels(_ x: Float) -> (value: Float, differential: (Float) -> Float) {
+    return (x, { $0 })
+  }
+  static func vjpResultLabels(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+    return (x, { $0 })
+  }
+
+  @differentiable(jvp: jvpResultLabels, vjp: vjpResultLabels)
+  func derivativeResultLabels(_ x: Float) -> Float {
+    return x
+  }
+  func jvpResultLabels(_ x: Float) -> (value: Float, differential: (Float) -> Float) {
+    return (x, { $0 })
+  }
+  func vjpResultLabels(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+    return (x, { $0 })
+  }
+}
+
 struct Tensor<Scalar> : AdditiveArithmetic {}
 extension Tensor : Differentiable where Scalar : Differentiable {
   typealias TangentVector = Tensor
@@ -544,7 +579,6 @@ extension FloatingPoint {
 }
 
 protocol MethodDiffReq {
-  // expected-error @+1 {{'vjpFoo' does not have expected type '<Self where Self : Differentiable, Self : MethodDiffReq> (Self) -> () -> (Self, (Self.CotangentVector) -> Self.CotangentVector)'}}
   @differentiable(wrt: self, vjp: vjpFoo where Self : Differentiable)
   func foo() -> Self
 }

branches/rxwei-patch-1/test/AutoDiff/differentiating_attr_type_checking.swift

@@ -280,3 +280,18 @@ extension InstanceMethodProto where Self : Differentiable {
     return (bar(), { _ in .zero })
   }
 }
+
+// Test consistent usages of `@differentiable` and `@differentiating` where
+// derivative functions are specified in both attributes.
+@differentiable(jvp: jvpConsistent, vjp: vjpConsistent)
+func consistentSpecifiedDerivatives(_ x: Float) -> Float {
+  return x
+}
+@differentiating(consistentSpecifiedDerivatives)
+func jvpConsistent(_ x: Float) -> (value: Float, differential: (Float) -> Float) {
+  return (x, { $0 })
+}
+@differentiating(consistentSpecifiedDerivatives(_:))
+func vjpConsistent(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+  return (x, { $0 })
+}