From 0ec663d09be0faf73d9bc544b78d010f2c228ad5 Mon Sep 17 00:00:00 2001
From: Victor Guerra <vm.guerramoran@criteo.com>
Date: Wed, 30 Oct 2019 12:49:03 +0100
Subject: [PATCH] [TF-812] Defines remaining derivatives for tgmath functions.

The following math functions are now differentiable:

* `remainder`
* `fmod`
* `ceil`
* `floor`
* `round`
* `trunc`

Resolves [TF-812](https://bugs.swift.org/projects/TF/issues/TF-812)
---
 stdlib/public/Platform/tgmath.swift.gyb | 86 +++++++++++++++++++++++++
 test/stdlib/tgmath.swift.gyb            | 10 +++
 2 files changed, 96 insertions(+)

diff --git a/stdlib/public/Platform/tgmath.swift.gyb b/stdlib/public/Platform/tgmath.swift.gyb
index 0d1526865353d..148c603602009 100644
--- a/stdlib/public/Platform/tgmath.swift.gyb
+++ b/stdlib/public/Platform/tgmath.swift.gyb
@@ -25,6 +25,7 @@ public func fabs<T: FloatingPoint>(_ x: T) -> T {
   vjp: _vjpSqrt
   where T : Differentiable & FloatingPoint, T == T.TangentVector
 )
+// SWIFT_ENABLE_TENSORFLOW END
 public func sqrt<T: FloatingPoint>(_ x: T) -> T {
   return x.squareRoot()
 }
@@ -36,36 +37,75 @@ public func sqrt<T: FloatingPoint>(_ x: T) -> T {
   vjp: _vjpFma
   where T : Differentiable & FloatingPoint, T == T.TangentVector
 )
+// SWIFT_ENABLE_TENSORFLOW END
 public func fma<T: FloatingPoint>(_ x: T, _ y: T, _ z: T) -> T {
   return z.addingProduct(x, y)
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  wrt: (x, y),
+  vjp: _vjpRemainder
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func remainder<T: FloatingPoint>(_ x: T, _ y: T) -> T {
   return x.remainder(dividingBy: y)
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  wrt: (x, y),
+  vjp: _vjpFmod
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func fmod<T: FloatingPoint>(_ x: T, _ y: T) -> T {
   return x.truncatingRemainder(dividingBy: y)
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  vjp: _vjpCeil
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func ceil<T: FloatingPoint>(_ x: T) -> T {
   return x.rounded(.up)
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  vjp: _vjpFloor
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func floor<T: FloatingPoint>(_ x: T) -> T {
   return x.rounded(.down)
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  vjp: _vjpRound
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func round<T: FloatingPoint>(_ x: T) -> T {
   return x.rounded()
 }
 
 @_transparent
+// SWIFT_ENABLE_TENSORFLOW
+@differentiable(
+  vjp: _vjpTrunc
+  where T : Differentiable & FloatingPoint, T == T.TangentVector
+)
+// SWIFT_ENABLE_TENSORFLOW END
 public func trunc<T: FloatingPoint>(_ x: T) -> T {
   return x.rounded(.towardZero)
 }
@@ -111,6 +151,51 @@ func _vjpFma<T: FloatingPoint & Differentiable> (
   return (fma(x, y, z), { v in (v * y, v * x, v) })
 }
 
+@usableFromInline
+func _vjpRemainder<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (T, (T) -> (T, T)) where T == T.TangentVector {
+  return (remainder(x, y), { v in (v, v) })
+}
+
+@usableFromInline
+func _vjpFmod<T: FloatingPoint & Differentiable> (
+  _ x: T,
+  _ y: T
+) -> (T, (T) -> (T, T)) where T == T.TangentVector {
+  return (fmod(x, y), { v in (v, v) })
+}
+
+@usableFromInline
+func _vjpCeil<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (T, (T) -> T) where T == T.TangentVector {
+  return (ceil(x), { v in 0 })
+}
+
+@usableFromInline
+func _vjpFloor<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (T, (T) -> T) where T == T.TangentVector {
+  return (floor(x), { v in 0 })
+}
+
+@usableFromInline
+func _vjpRound<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (T, (T) -> T) where T == T.TangentVector {
+  return (round(x), { v in 0 })
+}
+
+@usableFromInline
+func _vjpTrunc<T: FloatingPoint & Differentiable> (
+  _ x: T
+) -> (T, (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 +318,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
diff --git a/test/stdlib/tgmath.swift.gyb b/test/stdlib/tgmath.swift.gyb
index 65341df93ab14..f2eea08eecc74 100644
--- a/test/stdlib/tgmath.swift.gyb
+++ b/test/stdlib/tgmath.swift.gyb
@@ -273,6 +273,16 @@ 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)
+  let remainderGrad = gradient(at: 4.0 as ${T}, 5.0 as ${T}, in: { x, y in remainder(x, y) })
+  expectEqualWithTolerance(1.0, remainderGrad.0, ulps: 16)
+  expectEqualWithTolerance(1.0, remainderGrad.1, ulps: 16)
+  let fmodGrad = gradient(at: 4.0 as ${T}, 5.0 as ${T}, in: { x, y in fmod(x, y) })
+  expectEqualWithTolerance(1.0, fmodGrad.0, ulps: 16)
+  expectEqualWithTolerance(1.0, fmodGrad.1, 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)
 }
 %end