[AutoDiff][stdlib] Add JVPs to SIMDDifferentiation.swift.gyb (#32854)

Author: efremale

stdlib/public/Differentiation/SIMDDifferentiation.swift.gyb

@@ -59,6 +59,16 @@ where
       return zeros
     })
   }
+
+  @inlinable
+  @derivative(of: subscript(_:))
+  internal func _jvpSubscript(index: Int)
+    -> (value: Scalar, differential: (TangentVector) -> Scalar.TangentVector)
+  {
+    return (self[index], { v in
+      return .init(v[index])
+    })
+  }
 }
 
 %end
@@ -82,6 +92,18 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: +)
+  static func _jvpAdd(lhs: Self, rhs: Self)
+    -> (
+      value: Self, differential: (TangentVector, TangentVector) -> TangentVector
+    )
+  {
+    return (lhs + rhs, { ltan, rtan in
+      return ltan + rtan
+    })
+  }
+
   @inlinable
   @derivative(of: -)
   static func _vjpSubtract(lhs: Self, rhs: Self)
@@ -94,6 +116,18 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: -)
+  static func _jvpSubtract(lhs: Self, rhs: Self)
+    -> (
+      value: Self, differential: (TangentVector, TangentVector) -> TangentVector
+    )
+  {
+    return (lhs - rhs, { ltan, rtan in
+      return ltan - rtan
+    })
+  }
+
   @inlinable
   @derivative(of: -)
   static func _vjpNegate(rhs: Self)
@@ -103,6 +137,16 @@ where
       return -v
     })
   }
+
+  @inlinable
+  @derivative(of: -)
+  static func _jvpNegate(rhs: Self)
+    -> (value: Self, differential: (TangentVector) -> (TangentVector))
+  {
+    return (-rhs, { v in
+      return -v
+    })
+  }
 }
 
 extension SIMD
@@ -124,6 +168,18 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: *)
+  static func _jvpMultiply(lhs: Self, rhs: Self)
+    -> (
+      value: Self,  differential: (TangentVector, TangentVector) -> TangentVector
+    )
+  {
+    return (lhs * rhs, { ltan, rtan in
+      return lhs * ltan + rtan * rhs
+    })
+  }
+
   @inlinable
   @derivative(of: /)
   static func _vjpDivide(lhs: Self, rhs: Self)
@@ -135,6 +191,18 @@ where
       (v / rhs, -lhs / (rhs * rhs) * v)
     })
   }
+
+  @inlinable
+  @derivative(of: /)
+  static func _jvpDivide(lhs: Self, rhs: Self)
+    -> (
+      value: Self, differential: (TangentVector, TangentVector) -> TangentVector
+    )
+  {
+    return ( lhs / rhs, { ltan, rtan in
+      (ltan * rhs - lhs * rtan) / (rhs * rhs)
+    })
+  }
 }
 
 extension SIMD
@@ -156,6 +224,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: +)
+  static func _jvpAdd(lhs: Scalar, rhs: Self) -> (
+    value: Self,
+    differential: (Scalar.TangentVector, TangentVector) -> TangentVector
+  ) {
+    return (lhs + rhs, { ltan, rtan in
+      return ltan + rtan
+    })
+  }
+
   @inlinable
   @derivative(of: -)
   static func _vjpSubtract(lhs: Scalar, rhs: Self) -> (
@@ -167,6 +246,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: -)
+  static func _jvpSubtract(lhs: Scalar, rhs: Self) -> (
+    value: Self,
+    differential: (Scalar.TangentVector, TangentVector) -> TangentVector
+  ) {
+    return (lhs - rhs, { ltan, rtan in
+      return ltan - rtan
+    })
+  }
+
   @inlinable
   @derivative(of: +)
   static func _vjpAdd(lhs: Self, rhs: Scalar) -> (
@@ -178,6 +268,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: +)
+  static func _jvpAdd(lhs: Self, rhs: Scalar) -> (
+    value: Self,
+    differential: (TangentVector, Scalar.TangentVector) -> TangentVector
+  ) {
+    return (lhs + rhs, { ltan, rtan in
+      return ltan + rtan
+    })
+  }
+
   @inlinable
   @derivative(of: -)
   static func _vjpSubtract(lhs: Self, rhs: Scalar) -> (
@@ -188,6 +289,17 @@ where
       return (v, -v.sum())
     })
   }
+
+  @inlinable
+  @derivative(of: -)
+  static func _jvpSubtract(lhs: Self, rhs: Scalar) -> (
+    value: Self,
+    differential: (TangentVector, Scalar.TangentVector) -> TangentVector
+  ) {
+    return (lhs - rhs, { ltan, rtan in
+      return ltan - rtan
+    })
+  }
 }
 
 extension SIMD
@@ -209,6 +321,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: *)
+  static func _jvpMultiply(lhs: Self, rhs: Scalar) -> (
+    value: Self,
+    differential: (TangentVector, Scalar.TangentVector) -> TangentVector
+  ) {
+    return (lhs * rhs, { ltan, rtan in
+      return lhs * rtan + ltan * rhs
+    })
+  }
+
   @inlinable
   @derivative(of: /)
   static func _vjpDivide(lhs: Self, rhs: Scalar) -> (
@@ -220,6 +343,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: /)
+  static func _jvpDivide(lhs: Self, rhs: Scalar) -> (
+    value: Self,
+    differential: (TangentVector, Scalar.TangentVector) -> TangentVector
+  ) {
+    return (lhs / rhs, { ltan, rtan in
+      (ltan * rhs - lhs * rtan) / (rhs * rhs)
+    })
+  }
+
   @inlinable
   @derivative(of: *)
   static func _vjpMultiply(lhs: Scalar, rhs: Self) -> (
@@ -231,6 +365,17 @@ where
     })
   }
 
+  @inlinable
+  @derivative(of: *)
+  static func _jvpMultiply(lhs: Scalar, rhs: Self) -> (
+    value: Self,
+    differential: (Scalar.TangentVector, TangentVector) -> TangentVector
+  ) {
+    return (lhs * rhs, { ltan, rtan in
+      return lhs * rtan + ltan * rhs
+    })
+  }
+
   @inlinable
   @derivative(of: /)
   static func _vjpDivide(lhs: Scalar, rhs: Self) -> (
@@ -241,6 +386,17 @@ where
       ((v / rhs).sum(), -lhs / (rhs * rhs) * v)
     })
   }
+
+  @inlinable
+  @derivative(of: /)
+  static func _jvpDivide(lhs: Scalar, rhs: Self) -> (
+    value: Self,
+    differential: (Scalar.TangentVector, TangentVector) -> TangentVector
+  ) {
+    return (lhs / rhs, { ltan, rtan in
+      (ltan * rhs - lhs * rtan) / (rhs * rhs)
+    })
+  }
 }
 
 // FIXME(TF-1103): Derivative registration does not yet support
@@ -261,6 +417,14 @@ where
   ) {
     return (sum(), { v in Self(repeating: Scalar(v)) })
   }
+
+  @inlinable
+  @derivative(of: sum)
+  func _jvpSum() -> (
+    value: Scalar, differential (TangentVector) -> Scalar.TangentVector
+  ) {
+    return (sum(), { v in v.sum() }
+  }
 }
 */
 
@@ -279,4 +443,12 @@ where
   {
     return (Self(repeating: value), { v in v.sum() })
   }
+
+  @inlinable
+  @derivative(of: init(repeating:))
+  static func _jvpInit(repeating value: Scalar)
+    -> (value: Self, differential: (Scalar.TangentVector) -> TangentVector)
+  {
+    return (Self(repeating: value), { v in Self(repeating: v) })
+  }
 }