feat: add in existing code

JaapWijnen · JaapWijnen · commit 198dae635176 · 2024-11-25T15:13:55.000+01:00
This copies in parts of our existing Differentiation library code. No tests yet since those are written using XCTest will add those in a separate MR.
diff --git a/Sources/Differentiation/Array+Update.swift b/Sources/Differentiation/Array+Update.swift
@@ -0,0 +1,51 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+
+public extension Array where Element: Differentiable {
+    // Note: a compiler bug (SR-15530: https://bugs.swift.org/browse/SR-15530)
+    // causes the vjpUpdate functions to be associated with the wrong
+    // base functions unless you exactly align the function signatures in the
+    // @derivative(of:) attribute and make sure the @inlinable attribute is the
+    // same.
+
+    /// This function defines a derivative for AutoDiff to use when update() is called. It's not meant to be called directly in most
+    /// situations.
+    ///
+    /// - Parameters:
+    ///     - index: The index to update the value at.
+    ///     - newValue: The value to write.
+    /// - Returns: The object, plus the pullback.
+    @inlinable
+    @derivative(of: update(at:with:))
+    mutating func vjpUpdate(
+        at index: Int,
+        with newValue: Element
+    ) -> (value: Void, pullback: (inout TangentVector) -> (Element.TangentVector)) {
+        update(at: index, with: newValue)
+        let forwardCount = self.count
+        return ((), { tangentVector in
+            // manual zero tangent initialization
+            if tangentVector.base.count < forwardCount {
+                tangentVector.base = .init(repeating: .zero, count: forwardCount)
+            }
+            let dElement = tangentVector[index]
+            tangentVector.base[index] = .zero
+            return dElement
+        })
+    }
+}
+#endif
+
+public extension Array {
+    /// A functional version of `Array.subscript.modify`.
+    /// Differentiation does yet not support `Array.subscript.modify` because
+    /// it is a coroutine.
+    @inlinable
+    #if canImport(_Differentiation)
+    @differentiable(reverse where Element: Differentiable)
+    #endif
+    mutating func update(at index: Int, with newValue: Element) {
+        self[index] = newValue
+    }
+}
diff --git a/Sources/Differentiation/DerivativesOfNativeFunctions.swift b/Sources/Differentiation/DerivativesOfNativeFunctions.swift
@@ -0,0 +1,76 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+import Foundation
+
+// -------------------------------------------------------------------------
+// derivatives for native functions
+// -------------------------------------------------------------------------
+
+/// For min(): "Returns: The lesser of `x` and `y`. If `x` is equal to `y`, returns `x`."
+/// https://github.com/apple/swift/blob/main/stdlib/public/core/Algorithm.swift#L18
+@inlinable
+@derivative(of: min)
+public func minVJP<T: Comparable & Differentiable>(
+    _ lhs: T,
+    _ rhs: T
+) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+    func pullback(_ tangentVector: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+        if lhs <= rhs {
+            return (tangentVector, .zero)
+        }
+        else {
+            return (.zero, tangentVector)
+        }
+    }
+    return (value: min(lhs, rhs), pullback: pullback)
+}
+
+/// For max(): "Returns: The greater of `x` and `y`. If `x` is equal to `y`, returns `y`."
+/// https://github.com/apple/swift/blob/main/stdlib/public/core/Algorithm.swift#L52
+@inlinable
+@derivative(of: max)
+public func maxVJP<T: Comparable & Differentiable>(
+    _ lhs: T,
+    _ rhs: T
+) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+    func pullback(_ tangentVector: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+        if lhs < rhs {
+            return (.zero, tangentVector)
+        }
+        else {
+            return (tangentVector, .zero)
+        }
+    }
+    return (value: max(lhs, rhs), pullback: pullback)
+}
+
+/// To differentiate ``abs``
+@inlinable
+@derivative(of: abs)
+public func absVJP<T: Comparable & SignedNumeric & Differentiable>(_ value: T)
+    -> (value: T, pullback: (T.TangentVector) -> T.TangentVector)
+{
+    func pullback(_ tangentVector: T.TangentVector) -> T.TangentVector {
+        if value < 0 {
+            return .zero - tangentVector
+        }
+        else {
+            return tangentVector
+        }
+    }
+    return (value: abs(value), pullback: pullback)
+}
+
+/// Differentiation of ``atan2``
+@derivative(of: atan2(_:_:))
+public func vjpAtan2(
+    y: Double, x: Double
+) -> (value: Double, pullback: (Double) -> (Double, Double)) {
+    (
+        value: atan2(y, x),
+        pullback: { ($0 * x / (x * x + y * y), -$0 * y / (x * x + y * y)) }
+    )
+}
+
+#endif
diff --git a/Sources/Differentiation/Dictionary+Differentiation.swift b/Sources/Differentiation/Dictionary+Differentiation.swift
@@ -0,0 +1,114 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+
+// copied and modified from
+// https://github.com/borglab/SwiftFusion/blob/main/Sources/SwiftFusion/Core/Dictionary+Differentiable.swift
+// and
+// https://bugs.swift.org/browse/TF-1193
+
+/// This file makes `Dictionary` differentiable.
+///
+/// Note: This will eventually be moved into the Swift standard library. Once it is in the
+/// standard library, we can delete it from this repository.
+/// Implements the `Differentiable` requirements.
+extension Dictionary: Differentiable where Value: Differentiable {
+    public typealias TangentVector = [Key: Value.TangentVector]
+    public mutating func move(by direction: TangentVector) {
+        for (componentKey, componentDirection) in direction {
+            func fatalMissingComponent() -> Value {
+                preconditionFailure("missing component \(componentKey) in moved Dictionary")
+            }
+            self[componentKey, default: fatalMissingComponent()].move(by: componentDirection)
+        }
+    }
+
+    public var zeroTangentVectorInitializer: () -> TangentVector {
+        let listOfKeys = keys // capturing only what's needed, not the entire self, in order to not waste memory
+        func initializer() -> Self.TangentVector {
+            return listOfKeys.reduce(into: [Key: Value.TangentVector]()) { $0[$1] = Value.TangentVector.zero }
+        }
+        return initializer
+    }
+}
+
+/// Implements the `AdditiveArithmetic` requirements.
+extension Dictionary: AdditiveArithmetic where Value: AdditiveArithmetic {
+    public static func + (_ lhs: Self, _ rhs: Self) -> Self {
+        lhs.merging(rhs, uniquingKeysWith: +)
+    }
+
+    public static func - (_ lhs: Self, _ rhs: Self) -> Self {
+        lhs.merging(rhs.mapValues { .zero - $0 }, uniquingKeysWith: +)
+    }
+
+    public static var zero: Self { [:] }
+}
+
+// attempt to make builtin subscript differentiable:
+// https://bugs.swift.org/browse/TF-1193
+// https://github.com/apple/swift/pull/32614/
+// https://github.com/borglab/SwiftFusion/blob/main/Sources/SwiftFusion/Core/Dictionary+Differentiable.swift
+
+extension Dictionary where Value: Differentiable {
+    // get
+    // swiftformat:disable:next typeSugar
+    // periphery:ignore
+    @usableFromInline
+    @derivative(of: subscript(_:))
+    func vjpSubscriptGet(key: Key)
+        -> (value: Value?, pullback: (Optional<Value>.TangentVector) -> Dictionary<Key, Value>.TangentVector)
+    {
+        // When adding two dictionaries, nil values are equivalent to zeroes, so there is no need to manually zero-out
+        // every key's value. Instead, it is faster to create a dictionary with the single non-zero entry.
+        return (self[key], { tangentVector in
+            if let value = tangentVector.value {
+                return [key: value]
+            }
+            else {
+                return .zero
+            }
+        })
+    }
+}
+
+public extension Dictionary where Value: Differentiable {
+    // make a manual update(at: with:) since https://bugs.swift.org/browse/TF-1277 affects dictionary as well, making @derivative(of:
+    // subscript(_:).set) useless
+    /// manual update function replacing `subscript(_:).set` since that cannot be made differentiable (might now be possible)
+    @differentiable(reverse)
+    mutating func set(_ key: Key, to newValue: Value) {
+        self[key] = newValue
+    }
+
+    /// derivative of above set function. Ideally this would just be the derivative of `subscript(_:).set`
+    @derivative(of: set)
+    mutating func vjpUpdated(
+        _ key: Key,
+        to newValue: Value
+    ) -> (value: Void, pullback: (inout TangentVector) -> (Value.TangentVector)) {
+        set(key, to: newValue)
+
+        let forwardCount = count
+        let forwardKeys = keys // may be heavy to capture all of these, not sure how to do without them though
+
+        return ((), { tangentVector in
+            // manual zero tangent initialization
+            if tangentVector.count < forwardCount {
+                tangentVector = Self.TangentVector()
+                forwardKeys.forEach { tangentVector[$0] = .zero }
+            }
+
+            if let dElement = tangentVector[key] {
+                tangentVector[key] = .zero
+                return dElement
+            }
+            else { // should this fail?
+                tangentVector[key] = .zero
+                return .zero
+            }
+        })
+    }
+}
+
+#endif
diff --git a/Sources/Differentiation/Optional+DifferentiableMap.swift b/Sources/Differentiation/Optional+DifferentiableMap.swift
@@ -0,0 +1,55 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+
+// TODO: This will soon be upstreamed to the swift toolchain after which it will be removed from this package.
+
+extension Optional where Wrapped: Differentiable {
+    @inlinable
+    @differentiable(reverse, wrt: self)
+    public func differentiableMap<Result: Differentiable>(
+        _ body: @differentiable(reverse) (Wrapped) -> Result
+    ) -> Optional<Result> {
+        map(body)
+    }
+
+    @inlinable
+    @derivative(of: differentiableMap)
+    internal func _vjpDifferentiableMap<Result: Differentiable>(
+        _ body: @differentiable(reverse) (Wrapped) -> Result
+    ) -> (
+        value: Optional<Result>,
+        pullback: (Optional<Result>.TangentVector) -> Optional.TangentVector
+    ) {
+        let vwpb = self.map { valueWithPullback(at: $0, of: body) }
+        let bodyPullback = vwpb?.pullback
+
+        func pullback(_ vec: Optional<Result>.TangentVector) -> Optional.TangentVector {
+            guard let value = vec.value, let bodyPullback else { return .init(.none) }
+            return .init(.some(bodyPullback(value)))
+        }
+
+        return (value: vwpb?.value, pullback: pullback)
+    }
+
+    @inlinable
+    @derivative(of: differentiableMap)
+    internal func _jvpDifferentiableMap<Result: Differentiable>(
+        _ body: @differentiable(reverse) (Wrapped) -> Result
+    ) -> (
+        value: Optional<Result>,
+        differential: (Optional.TangentVector) -> Optional<Result>.TangentVector
+    ) {
+        let vwpb = self.map { valueWithDifferential(at: $0, of: body) }
+        let bodyDifferential = vwpb?.differential
+
+        func differential(_ vec: Optional.TangentVector) -> Optional<Result>.TangentVector {
+            guard let value = vec.value, let bodyDifferential else { return .init(.none) }
+            return .init(bodyDifferential(value))
+        }
+
+        return (value: vwpb?.value, differential: differential)
+    }
+}
+
+#endif
diff --git a/Sources/Differentiation/Sequence+MinMaxVJPs.swift b/Sources/Differentiation/Sequence+MinMaxVJPs.swift
@@ -0,0 +1,76 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+
+public extension Sequence where
+    Self: Collection, // we constrain to conform to collection cause otherwise we can't access any values by index
+    Self: Differentiable,
+    Self.TangentVector: RangeReplaceableCollection, // we constrain the tangentvector to be able to create a value and write to it
+    Self.TangentVector.Element == Element.TangentVector,
+    Element: Differentiable,
+    Element: Comparable
+{
+    // Match Self.Index  with Self.TangentVector index so we can use them across both types.
+    // The reason we are doing the where clause here rather than at the extension declaration
+    // level is because of the DocC crash: https://github.com/swiftlang/swift/issues/75258
+    /// To differentiate ``Swift/Sequence/max``
+    @derivative(of: max)
+    func vjpMax() -> (
+        value: Element?,
+        pullback: (Element?.TangentVector) -> (Self.TangentVector)
+    ) where Self.Index == Self.TangentVector.Index {
+        let index = withoutDerivative(at: self.indices.max { self[$0] < self[$1] }) // we grab the index of the element with the max value
+        return (
+            value: index.map { self[$0] }, // if the index is nil, we return nil otherwise we grab the value at the index
+            pullback: { vector in
+                var dSelf = Self
+                    .TangentVector(
+                        repeating: .zero,
+                        count: self
+                            .count
+                    ) // we create a zero tangentvector we need `RangeReplaceableCollection` conformance in order to do this
+                if let vectorValue = vector.value,
+                   let index = index
+                {
+                    // if an index was found and our tangentvector's value is non nil we set the value at index of our tangentvector to the
+                    // provided tangentvector value
+                    dSelf
+                        .replaceSubrange(
+                            index ..< dSelf.index(after: index),
+                            with: [vectorValue]
+                        ) // we use `RangeReplaceableCollection`'s method here in order to not have to also constrain our TangentVector to
+                    // `MutableCollection`
+                }
+                return dSelf // return the tangentvector
+            }
+        )
+    }
+
+    // Match Self.Index  with Self.TangentVector index so we can use them across both types.
+    // The reason we are doing the where clause here rather than at the extension declaration
+    // level is because of the DocC crash: https://github.com/swiftlang/swift/issues/75258
+    /// To differentiate ``Swift/Sequence/min``
+    @derivative(of: min)
+    func vjpMin() -> (
+        value: Element?,
+        pullback: (Element?.TangentVector) -> (Self.TangentVector)
+    ) where Self.Index == Self.TangentVector.Index {
+        let index = withoutDerivative(at: self.indices.min { self[$0] < self[$1] }) // we grab the index of the element with the max value
+        return (
+            value: index.map { self[$0] }, // if the index is nil, we return nil otherwise we grab the value at the index
+            pullback: { vector in
+                var dSelf = Self.TangentVector(repeating: .zero, count: self.count) // we create a zero tangentvector
+                if let vectorValue = vector.value,
+                   let index = index
+                {
+                    // if an index was found and our tangentvector's value is non nil we set the value at index of our tangentvector to the
+                    // provided tangentvector value
+                    dSelf.replaceSubrange(index ..< dSelf.index(after: index), with: [vectorValue])
+                }
+                return dSelf // return the tangentvector
+            }
+        )
+    }
+}
+
+#endif
