[AutoDiff] Defines remaining derivatives for tgmath functions. (#28108)

Author: vguerra

stdlib/public/Platform/tgmath.swift.gyb

@@ -20,22 +20,11 @@ public func fabs<T: FloatingPoint>(_ x: T) -> T {
 }
 
 @_transparent
-// SWIFT_ENABLE_TENSORFLOW
-@differentiable(
-  vjp: _vjpSqrt
-  where T : Differentiable & FloatingPoint, T == T.TangentVector
-)
 public func sqrt<T: FloatingPoint>(_ x: T) -> T {
   return x.squareRoot()
 }
 
 @_transparent
-// SWIFT_ENABLE_TENSORFLOW
-@differentiable(
-  wrt: (x, y, z),
-  vjp: _vjpFma
-  where T : Differentiable & FloatingPoint, T == T.TangentVector
-)
 public func fma<T: FloatingPoint>(_ x: T, _ y: T, _ z: T) -> T {
   return z.addingProduct(x, y)
 }
@@ -95,22 +84,75 @@ public func frexp<T: BinaryFloatingPoint>(_ x: T) -> (T, Int) {
 
 // SWIFT_ENABLE_TENSORFLOW
 @usableFromInline
+@differentiating(sqrt)
 func _vjpSqrt<T: FloatingPoint & Differentiable> (
   _ x: T
-) -> (T, (T) -> T) where T == T.TangentVector {
-  let value = x.squareRoot()
+) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+  let value = sqrt(x)
   return (value, { v in v / (2 * value) })
 }
 
 @usableFromInline
+@differentiating(fma)
 func _vjpFma<T: FloatingPoint & Differentiable> (
   _ x: T,
   _ y: T,
   _ z: T
-) -> (T, (T) -> (T, T, T)) where T == T.TangentVector {
+) -> (value: T, pullback: (T) -> (T, T, T)) where T == T.TangentVector {
   return (fma(x, y, z), { v in (v * y, v * x, v) })
 }
 
+@usableFromInline
+@differentiating(remainder)
+func _vjpRemainder<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (value: T, pullback: (T) -> (T, T)) where T == T.TangentVector {
+  return (remainder(x, y), { v in (v, -v * ((x / y).rounded(.toNearestOrEven))) })
+}
+
+@usableFromInline
+@differentiating(fmod)
+func _vjpFmod<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (value: T, pullback: (T) -> (T, T)) where T == T.TangentVector {
+  return (fmod(x, y), { v in (v, -v * ((x / y).rounded(.towardZero))) })
+}
+
+@usableFromInline
+@differentiating(ceil)
+func _vjpCeil<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+  return (ceil(x), { v in 0 })
+}
+
+@usableFromInline
+@differentiating(floor)
+func _vjpFloor<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+  return (floor(x), { v in 0 })
+}
+
+@usableFromInline
+@differentiating(round)
+func _vjpRound<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+  return (round(x), { v in 0 })
+}
+
+@usableFromInline
+@differentiating(trunc)
+func _vjpTrunc<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (value: T, pullback: (T) -> T) where T == T.TangentVector {
+  return (trunc(x), { v in 0 })
+}
+// SWIFT_ENABLE_TENSORFLOW END
+
 %for T in ['Float','Double']:
 @available(swift, deprecated: 4.2, renamed: "scalbn")
 @_transparent
@@ -233,6 +275,7 @@ func _vjpErf(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
 func _vjpErfc(_ x: ${T}) -> (${T}, (${T}) -> ${T}) {
   return (erfc(x), { v in v * -${T}(M_2_SQRTPI) * exp(-x * x) })
 }
+// SWIFT_ENABLE_TENSORFLOW END
 %   if T == 'Float80':
 #endif
 %   end

test/stdlib/tgmath.swift.gyb

@@ -53,6 +53,19 @@ func expectEqualWithTolerance<T>(_ expected: TestLiteralType, _ actual: T,
              file: file, line: line)
 }
 
+func checkGradient<T: BinaryFloatingPoint & Differentiable>(
+  _ f: @differentiable (T, T) -> T,
+  _ x: T,
+  _ y: T)
+where T == T.TangentVector {
+  let eps = T(0.01)
+  let grad = gradient(at: x, y, in: f)
+  let dfdx = (f(x + eps, y) - f(x, y)) / eps
+  let dfdy = (f(x, y + eps) - f(x, y)) / eps
+  expectEqualWithTolerance(TestLiteralType(dfdx), grad.0, ulps: 192)
+  expectEqualWithTolerance(TestLiteralType(dfdy), grad.1, ulps: 192)
+}
+
 %{
 unary = [
   'acos', 'asin', 'atan',
@@ -273,6 +286,21 @@ MathTests.test("gradient_${T}") {
   expectEqualWithTolerance(5.0, fmaGrad.0, ulps: 16)
   expectEqualWithTolerance(4.0, fmaGrad.1, ulps: 16)
   expectEqualWithTolerance(1.0, fmaGrad.2, ulps: 16)
+  expectEqualWithTolerance(0.0, gradient(at: 2.0 as ${T}, in: { ceil($0) }), ulps: 16)
+  expectEqualWithTolerance(0.0, gradient(at: 2.0 as ${T}, in: { floor($0) }), ulps: 16)
+  expectEqualWithTolerance(0.0, gradient(at: 2.0 as ${T}, in: { round($0) }), ulps: 16)
+  expectEqualWithTolerance(0.0, gradient(at: 2.0 as ${T}, in: { trunc($0) }), ulps: 16)
+  for a in -10...10 {
+    let x = ${T}(a)
+    for b in -10...10 {
+      let y = ${T}(b)
+      guard b != 0 && remainder(x, y).sign == remainder(x + ${T}(0.001), y).sign &&
+        remainder(x, y).sign == remainder(x, y + ${T}(0.001)).sign
+        else { continue }
+      checkGradient({ remainder($0, $1) }, x, y)
+      checkGradient({ fmod($0, $1) }, x, y)
+    }
+  }
 }
 %end