[AutoDiff] Change Differentiable.moved(along:) to move(along:). (#25437)

Change `Differentiable.moved(along:)` to `mutating func move(along:)`.
This is important for upcoming class differentiation support.

Update `Differentiable` derived conformances (logic and diagnostics).
Update tests.

Author: dan-zheng

include/swift/AST/DiagnosticsSema.def

@@ -2538,12 +2538,15 @@ ERROR(broken_differentiable_requirement,none,
       "Differentiable protocol is broken: unexpected requirement", ())
 WARNING(differentiable_nondiff_type_implicit_noderivative_fixit,none,
       "stored property %0 has no derivative because it does not conform to "
-      "'Differentiable'; add '@noDerivative' to make it explicit",
-      (Identifier))
-WARNING(differentiable_constant_property_implicit_noderivative_fixit,none,
-      "'let' properties with a default value do not have a derivative; add "
-      "'@noDerivative' to make it explicit, or change it to 'var' to allow "
-      "derivatives", ())
+      "'Differentiable'; add an explicit '@noDerivative' attribute"
+      "%select{|, or conform %1 to 'AdditiveArithmetic'}2",
+      (Identifier, Identifier, bool))
+WARNING(differentiable_let_property_implicit_noderivative_fixit,none,
+      "synthesis of the 'Differentiable.move(along:)' requirement for %1 "
+      "requires all stored properties to be mutable; use 'var' instead, or add "
+      "an explicit '@noDerivative' attribute"
+      "%select{|, or conform %1 to 'AdditiveArithmetic'}2",
+      (Identifier, Identifier, bool))
 
 NOTE(codable_extraneous_codingkey_case_here,none,
      "CodingKey case %0 does not match any stored properties", (Identifier))

include/swift/AST/KnownIdentifiers.def

@@ -140,7 +140,7 @@ IDENTIFIER(scaled)
 IDENTIFIER(AllDifferentiableVariables)
 IDENTIFIER(TangentVector)
 IDENTIFIER(allDifferentiableVariables)
-IDENTIFIER(moved)
+IDENTIFIER(move)
 
 // Kinds of layout constraints
 IDENTIFIER_WITH_NAME(UnknownLayout, "_UnknownLayout")

lib/Sema/DerivedConformanceDifferentiable.cpp

@@ -315,9 +315,9 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,
     }
 
     // SWIFT_ENABLE_TENSORFLOW
-    // Differentiable.moved(along:)
+    // Differentiable.move(along:)
     if (name.isCompoundName() &&
-        name.getBaseName() == ctx.Id_moved) {
+        name.getBaseName() == ctx.Id_move) {
       auto argumentNames = name.getArgumentNames();
       if (argumentNames.size() == 1 &&
           argumentNames[0] == ctx.getIdentifier("along")) {

lib/Sema/DerivedConformances.cpp

@@ -1973,13 +1973,14 @@ extension Array where Element : Differentiable {
       }
     }
 
-    public func moved(along direction: TangentVector) -> DifferentiableView {
+    public mutating func move(along direction: TangentVector) {
       precondition(
         base.count == direction.base.count,
         "cannot move Array.DifferentiableView with count \(base.count) along " +
           "direction with different count \(direction.base.count)")
-      return DifferentiableView(
-        zip(base, direction.base).map { $0.moved(along: $1) })
+      for i in base.indices {
+        base[i].move(along: direction.base[i])
+      }
     }
   }
 }
@@ -2072,8 +2073,10 @@ extension Array : Differentiable where Element : Differentiable {
     }
   }
 
-  public func moved(along direction: TangentVector) -> Array {
-    return DifferentiableView(self).moved(along: direction).base
+  public mutating func move(along direction: TangentVector) {
+    var view = DifferentiableView(self)
+    view.move(along: direction)
+    self = view.base
   }
 }
 

stdlib/public/core/Array.swift

@@ -65,10 +65,9 @@ public protocol Differentiable {
   /// All differentiable variables of this value.
   var allDifferentiableVariables: AllDifferentiableVariables { get set }
 
-  /// Returns `self` moved along the value space towards the given tangent
-  /// vector. In Riemannian geometry (mathematics), this represents an
-  /// exponential map.
-  func moved(along direction: TangentVector) -> Self
+  /// Moves `self` along the value space towards the given tangent vector. In
+  /// Riemannian geometry (mathematics), this represents an exponential map.
+  mutating func move(along direction: TangentVector)
 
   @available(*, deprecated,
              message: "'CotangentVector' is now equal to 'TangentVector' and will be removed")
@@ -82,13 +81,9 @@ public extension Differentiable where AllDifferentiableVariables == Self {
   }
 }
 
-// FIXME: The `Self : AdditiveArithmetic` constraint should be implied by
-// `TangentVector == Self`, but the type checker errors out when it does not
-// exist.
-public extension Differentiable
-  where TangentVector == Self, Self : AdditiveArithmetic {
-  func moved(along direction: TangentVector) -> Self {
-    return self + direction
+public extension Differentiable where TangentVector == Self {
+  mutating func move(along direction: TangentVector) {
+    self += direction
   }
 }
 
@@ -451,7 +446,7 @@ internal protocol _AnyDerivativeBox {
 
   // `Differentiable` requirements.
   var _allDifferentiableVariables: _AnyDerivativeBox { get }
-  func _moved(along direction: _AnyDerivativeBox) -> _AnyDerivativeBox
+  mutating func _move(along direction: _AnyDerivativeBox)
 
   /// The underlying base value, type-erased to `Any`.
   var _typeErasedBase: Any { get }
@@ -555,18 +550,17 @@ internal struct _ConcreteDerivativeBox<T> : _AnyDerivativeBox
     return _ConcreteDerivativeBox(_base.allDifferentiableVariables)
   }
 
-  func _moved(along direction: _AnyDerivativeBox) -> _AnyDerivativeBox {
-    if _isOpaqueZero() {
-      return direction
-    }
+  mutating func _move(along direction: _AnyDerivativeBox) {
     if direction._isOpaqueZero() {
-      return self
+      return
     }
+    // The case where `self._isOpaqueZero()` returns true is handled in
+    // `AnyDerivative.move(along:)`.
     guard let directionBase =
       direction._unboxed(to: T.TangentVector.self) else {
       _derivativeTypeMismatch(T.self, type(of: direction._typeErasedBase))
     }
-    return _ConcreteDerivativeBox<T>(_base.moved(along: directionBase))
+    _base.move(along: directionBase)
   }
 }
 
@@ -661,7 +655,11 @@ public struct AnyDerivative : Differentiable & AdditiveArithmetic {
     get { return AnyDerivative(_box: _box._allDifferentiableVariables) }
     // set { _box._allDifferentiableVariables = newValue._box }
   }
-  public func moved(along direction: TangentVector) -> AnyDerivative {
-    return AnyDerivative(_box: _box._moved(along: direction._box))
+  public mutating func move(along direction: TangentVector) {
+    if _box._isOpaqueZero() {
+      _box = direction._box
+      return
+    }
+    _box._move(along: direction._box)
   }
 }

stdlib/public/core/AutoDiff.swift

@@ -1887,6 +1887,10 @@ extension ${Self} : VectorProtocol {
 extension ${Self} : Differentiable {
   public typealias TangentVector = ${Self}
   public typealias AllDifferentiableVariables = ${Self}
+
+  public mutating func move(along direction: TangentVector) {
+    self += direction
+  }
 }
 
 //===----------------------------------------------------------------------===//

stdlib/public/core/FloatingPointTypes.swift.gyb

@@ -6,10 +6,19 @@ import StdlibUnittest
 var AnyDerivativeTests = TestSuite("AnyDerivative")
 
 struct Vector : Differentiable {
-  let x, y: Float
+  var x, y: Float
 }
 struct Generic<T: Differentiable> : Differentiable {
-  let x: T
+  var x: T
+}
+
+extension AnyDerivative {
+  // This exists only to faciliate testing.
+  func moved(along direction: TangentVector) -> Self {
+    var result = self
+    result.move(along: direction)
+    return result
+  }
 }
 
 AnyDerivativeTests.test("Vector") {

test/AutoDiff/anyderivative.swift

@@ -29,8 +29,8 @@ ArrayAutodiffTests.test("ArraySubscript") {
 
 ArrayAutodiffTests.test("ArrayConcat") {
   struct TwoArrays : Differentiable {
-    let a: [Float]
-    let b: [Float]
+    var a: [Float]
+    var b: [Float]
   }
 
   func sumFirstThreeConcatted(_ arrs: TwoArrays) -> Float {

test/AutoDiff/array.swift

@@ -25,10 +25,10 @@ _ = gradient(at: 0, in: one_to_one_0) // okay!
 //===----------------------------------------------------------------------===//
 
 struct S {
-  let p: Float
+  var p: Float
 }
-
 extension S : Differentiable, VectorProtocol {
+  // Test custom `TangentVector` type with non-matching stored property name.
   struct TangentVector: Differentiable, VectorProtocol {
     var dp: Float
   }
@@ -39,8 +39,8 @@ extension S : Differentiable, VectorProtocol {
   static func - (lhs: S, rhs: S) -> S { return S(p: lhs.p - rhs.p) }
   static func * (lhs: Float, rhs: S) -> S { return S(p: lhs * rhs.p) }
 
-  func moved(along direction: TangentVector) -> S {
-    return S(p: p + direction.dp)
+  mutating func move(along direction: TangentVector) {
+    p.move(along: direction.dp)
   }
 }
 

test/AutoDiff/autodiff_diagnostics.swift

@@ -77,12 +77,12 @@ struct GenericTanMember<T : Differentiable> : Differentiable, AdditiveArithmetic
 // CHECK-AST:   @_implements(Equatable, ==(_:_:)) internal static func __derived_struct_equals(_ a: GenericTanMember<T>, _ b: GenericTanMember<T>) -> Bool
 
 public struct ConditionallyDifferentiable<T> {
-  public let x: T
+  public var x: T
 }
 extension ConditionallyDifferentiable : Differentiable where T : Differentiable {}
 
 // CHECK-AST-LABEL: public struct ConditionallyDifferentiable<T> {
 // CHECK-AST:         @differentiable(wrt: self where T : Differentiable)
-// CHECK-AST:         public let x: T
+// CHECK-AST:         public var x: T
 // CHECK-AST:         internal init(x: T)
 // CHECK-AST:       }

test/AutoDiff/derived_differentiable_properties.swift

@@ -548,11 +548,7 @@ struct ResultLabelTest {
 }
 
 struct Tensor<Scalar> : AdditiveArithmetic {}
-extension Tensor : Differentiable where Scalar : Differentiable {
-  typealias TangentVector = Tensor
-  typealias AllDifferentiableVariables = Tensor
-  func moved(along direction: Tensor) -> Tensor { return self }
-}
+extension Tensor : Differentiable where Scalar : Differentiable {}
 @differentiable(where Scalar : Differentiable)
 func where2<Scalar : Numeric>(x: Tensor<Scalar>) -> Tensor<Scalar> {
   return x

test/AutoDiff/differentiable_attr_type_checking.swift

@@ -20,16 +20,15 @@ struct Space {
   /// `x` is a computed property with a custom vjp.
   var x: Float {
     @differentiable(vjp: vjpX)
-    get {
-      return storedX
-    }
+    get { storedX }
+    set { storedX = newValue }
   }
 
   func vjpX() -> (Float, (Float) -> TangentSpace) {
     return (x, { v in TangentSpace(x: v, y: 0) } )
   }
 
-  private let storedX: Float
+  private var storedX: Float
 
   @differentiable
   var y: Float
@@ -42,8 +41,9 @@ struct Space {
 
 extension Space : Differentiable {
   typealias TangentVector = TangentSpace
-  func moved(along: TangentSpace) -> Space {
-    return Space(x: x + along.x, y: y + along.y)
+  mutating func move(along direction: TangentSpace) {
+    x.move(along: direction.x)
+    y.move(along: direction.y)
   }
 }
 
@@ -106,13 +106,14 @@ extension ProductSpaceOtherTangentTangentSpace : Differentiable {
 }
 
 struct ProductSpaceOtherTangent {
-  let x, y: Float
+  var x, y: Float
 }
 
 extension ProductSpaceOtherTangent : Differentiable {
   typealias TangentVector = ProductSpaceOtherTangentTangentSpace
-  func moved(along: ProductSpaceOtherTangentTangentSpace) -> ProductSpaceOtherTangent {
-    return ProductSpaceOtherTangent(x: x + along.x, y: y + along.y)
+  mutating func move(along direction: ProductSpaceOtherTangentTangentSpace) {
+    x.move(along: direction.x)
+    y.move(along: direction.y)
   }
 }