diff --git a/stdlib/public/Platform/tgmath.swift.gyb b/stdlib/public/Platform/tgmath.swift.gyb index 0d1526865353d..f2e4c1eace2a6 100644 --- a/stdlib/public/Platform/tgmath.swift.gyb +++ b/stdlib/public/Platform/tgmath.swift.gyb @@ -20,22 +20,11 @@ public func fabs(_ x: T) -> T { } @_transparent -// SWIFT_ENABLE_TENSORFLOW -@differentiable( - vjp: _vjpSqrt - where T : Differentiable & FloatingPoint, T == T.TangentVector -) public func sqrt(_ 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(_ x: T, _ y: T, _ z: T) -> T { return z.addingProduct(x, y) } @@ -95,22 +84,75 @@ public func frexp(_ x: T) -> (T, Int) { // SWIFT_ENABLE_TENSORFLOW @usableFromInline +@differentiating(sqrt) func _vjpSqrt ( _ 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 ( _ 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 ( + _ 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 ( + _ 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 ( + _ x: T +) -> (value: T, pullback: (T) -> T) where T == T.TangentVector { + return (ceil(x), { v in 0 }) +} + +@usableFromInline +@differentiating(floor) +func _vjpFloor ( + _ x: T +) -> (value: T, pullback: (T) -> T) where T == T.TangentVector { + return (floor(x), { v in 0 }) +} + +@usableFromInline +@differentiating(round) +func _vjpRound ( + _ x: T +) -> (value: T, pullback: (T) -> T) where T == T.TangentVector { + return (round(x), { v in 0 }) +} + +@usableFromInline +@differentiating(trunc) +func _vjpTrunc ( + _ 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 diff --git a/test/stdlib/tgmath.swift.gyb b/test/stdlib/tgmath.swift.gyb index 65341df93ab14..47cb8e13a6961 100644 --- a/test/stdlib/tgmath.swift.gyb +++ b/test/stdlib/tgmath.swift.gyb @@ -53,6 +53,19 @@ func expectEqualWithTolerance(_ expected: TestLiteralType, _ actual: T, file: file, line: line) } +func checkGradient( + _ 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