Adding derivatives for generic functions: sqrt and fma.

with respective tests.

Author: vguerra

stdlib/public/Platform/tgmath.swift.gyb

@@ -20,11 +20,20 @@ public func fabs<T: FloatingPoint>(_ x: T) -> T {
 }
 
 @_transparent
+@differentiable(
+  vjp: _vjpSqrt
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
 public func sqrt<T: FloatingPoint>(_ x: T) -> T {
   return x.squareRoot()
 }
 
 @_transparent
+@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)
 }
@@ -82,6 +91,25 @@ public func frexp<T: BinaryFloatingPoint>(_ x: T) -> (T, Int) {
   return (x.significand / 2, Int(x.exponent + 1))
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+@usableFromInline
+func _vjpSqrt<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (T, (T) -> T) where T == T.TangentVector {
+  let value = x.squareRoot()
+  return (value, {v in (1 / 2) * ( 1 / value) * v})
+}
+
+@usableFromInline
+func _vjpFma<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T,
+  _ z: T
+) -> (T, (T) -> (T, T, T)) where T == T.TangentVector {
+  return (fma(x, y, z),
+    {v in return (v * y, v * x, v)})
+}
+
 %for T in ['Float','Double']:
 @available(swift, deprecated: 4.2, renamed: "scalbn")
 @_transparent

test/stdlib/tgmath.swift.gyb

@@ -268,6 +268,11 @@ MathTests.test("gradient_${T}") {
   expectEqualWithTolerance(1.3333333333333333334, gradient(at: 0.5 as ${T}, in: atanh), ulps: 16)
   expectEqualWithTolerance(0.020666985354092053575, gradient(at: 2.0 as ${T}, in: erf), ulps: 16)
   expectEqualWithTolerance(-0.020666985354092053575, gradient(at: 2.0 as ${T}, in: erfc), ulps: 16)
+  expectEqualWithTolerance(0.35355339059327376222, gradient(at: 2.0 as ${T}, in: {x in sqrt(x)}), ulps: 16)
+  let fma_grad = gradient(at: 4.0 as ${T}, 5.0 as ${T}, 6.0 as ${T}, in: {x, y, z in fma(x, y, z)})
+  expectEqualWithTolerance(5.0, fma_grad.0, ulps: 16)
+  expectEqualWithTolerance(4.0, fma_grad.1, ulps: 16)
+  expectEqualWithTolerance(1.0, fma_grad.2, ulps: 16)
 }
 %end