feature: Make OrderedDictionary.Values conform to Differentiable

Author: JaapWijnen

Sources/OrderedCollectionsDifferentiable/OrderedDictionary+Differentiable.swift

@@ -5,12 +5,12 @@ import _Differentiation
 extension OrderedDictionary: @retroactive Differentiable where Value: Differentiable {
     public typealias TangentVector = OrderedDictionary<Key, Value.TangentVector>
 
-    public mutating func move(by direction: TangentVector) {
-        for (componentKey, componentDirection) in direction {
+    public mutating func move(by offset: TangentVector) {
+        for (key, tangentValue) in offset {
             func fatalMissingComponent() -> Value {
-                preconditionFailure("missing component \(componentKey) in moved OrderedDictionary")
+                preconditionFailure("missing entry for key \(key) in moved OrderedDictionary")
             }
-            self[componentKey, default: fatalMissingComponent()].move(by: componentDirection)
+            self[key, default: fatalMissingComponent()].move(by: tangentValue)
         }
     }
 }
@@ -33,9 +33,10 @@ extension OrderedDictionary where Value: Differentiable {
     /// differentiable
     @inlinable
     @derivative(of: subscript(_:))
-    func _vjpSubscript(key: Key)
-        -> (value: Value?, pullback: (Optional<Value>.TangentVector) -> OrderedDictionary<Key, Value>.TangentVector)
-    {
+    func _vjpSubscript(key: Key) -> (
+        value: Value?,
+        pullback: (Optional<Value>.TangentVector) -> OrderedDictionary<Key, Value>.TangentVector
+    ) {
         let keys = self.keys
         // 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.
@@ -58,8 +59,26 @@ extension OrderedDictionary where Value: Differentiable {
             }
         )
     }
+
+    @derivative(of: values)
+    @inlinable
+    @inline(__always)
+    public func _vjpValues() -> (value: Values, pullback: (Values.TangentVector) -> OrderedDictionary<Key, Value>.TangentVector) {
+        let keys = self.keys
+        return (
+            value: self.values,
+            pullback: { v in
+                var dict = OrderedDictionary<Key, Value>.TangentVector()
+                dict.reserveCapacity(keys.count)
+                for (key, tangentValue) in zip(keys, v.base) {
+                    dict[key] = tangentValue
+                }
+                return dict
+            }
+        )
+    }
 }
 
-// TODO: make `OrderedDictionary.Values` and `OrderedDictionary.Elements` differentiable
+// TODO: make `OrderedDictionary.Elements` differentiable
 
 #endif

Sources/OrderedCollectionsDifferentiable/OrderedDictionaryValues+Differentiable.swift

@@ -0,0 +1,48 @@
+#if canImport(_Differentiation)
+
+import _Differentiation
+
+extension OrderedDictionary.Values: @retroactive Differentiable where Value: Differentiable {
+    public typealias TangentVector = Array<Value.TangentVector>.TangentVector
+
+    public mutating func move(by offset: OrderedDictionary<Key, Value>.Values.TangentVector) {
+        for (i, j) in zip(self.indices, offset.base.indices) {
+            self[i].move(by: offset.base[j])
+        }
+    }
+}
+
+extension OrderedDictionary.Values where Value: Differentiable {
+    /// Defines a derivative for `OrderedDictionary`s subscript getter enabling calls like `var value = dictionary[key]` to be
+    /// differentiable
+    @inlinable
+    @derivative(of: subscript(_:))
+    func _vjpSubscript(position: Int) -> (
+        value: Value,
+        pullback: (Value.TangentVector) -> OrderedDictionary<Key, Value>.Values.TangentVector
+    ) {
+        let count = self.count
+        return (
+            value: self[position],
+            pullback: { tangentVector in
+                var vector = Array<Value.TangentVector>.TangentVector(.init(repeating: .zero, count: count))
+                vector.base[position] = tangentVector
+                return vector
+            }
+        )
+    }
+
+    @derivative(of: elements)
+    @inlinable
+    func _vjpElements() -> (
+        value: Array<Value>,
+        pullback: (Array<Value>.TangentVector) -> OrderedDictionary.Values.TangentVector
+    ) {
+        (
+            value: self.elements,
+            pullback: { v in v }
+        )
+    }
+}
+
+#endif

Tests/OrderedCollectionsDifferentiableTests/OrderedDictionary+DifferentiableTests.swift

@@ -3,8 +3,8 @@
 import OrderedCollectionsDifferentiable
 import Testing
 
-@Suite("Dictionary+Differentiation")
-struct DictionaryDifferentiationTests {
+@Suite("OrderedDictionary+Differentiation")
+struct OrderedDictionaryDifferentiationTests {
     @Test
     func testSubscriptGet() throws {
         let dictionary: OrderedDictionary<String, Double> = ["a": 3, "b": 7]

Tests/OrderedCollectionsDifferentiableTests/OrderedDictionaryValues+DifferentiableTests.swift

@@ -0,0 +1,59 @@
+#if canImport(_Differentiation)
+
+import OrderedCollectionsDifferentiable
+import Testing
+
+@Suite("OrderedDictionary.Values+Differentiable")
+struct OrderedDictionaryValuesTests {
+    @Test
+    func dictionaryValuesMemberTest() {
+        @differentiable(reverse)
+        func values(dict: OrderedDictionary<String, Double>) -> OrderedDictionary<String, Double>.Values {
+            dict.values
+        }
+
+        let dict: OrderedDictionary<String, Double> = ["a": 1.0, "b": 2.0]
+        let vwpb = valueWithPullback(at: dict, of: values)
+
+        #expect(vwpb.value == dict.values)
+        let pullback = vwpb.pullback([0.0, 1.0])
+
+        #expect(pullback == ["a": 0.0, "b": 1.0])
+    }
+
+    @Test
+    func dictionaryValuesMemberElementsTest() {
+        @differentiable(reverse)
+        func elements(dict: OrderedDictionary<String, Double>) -> [Double] {
+            dict.values.elements
+        }
+
+        let dict: OrderedDictionary<String, Double> = ["a": 1.0, "b": 2.0]
+        let vwpb = valueWithPullback(at: dict, of: elements)
+
+        #expect(vwpb.value == [1.0, 2.0])
+        let pullback = vwpb.pullback([0.0, 1.0])
+
+        #expect(pullback == ["a": 0.0, "b": 1.0])
+    }
+
+    @Test
+    func dictionaryValuesMemberSubscriptTest() {
+        @differentiable(reverse)
+        func values(dict: OrderedDictionary<String, Double>, index: Int) -> Double {
+            dict.values[index]
+        }
+
+        let dict: OrderedDictionary<String, Double> = ["a": 1.0, "b": 2.0]
+        let vwpb = valueWithPullback(at: dict, of: { dict in
+            values(dict: dict, index: 1)
+        })
+
+        #expect(vwpb.value == 2.0)
+        let pullback = vwpb.pullback(1.0)
+
+        #expect(pullback == ["a": 0.0, "b": 1.0])
+    }
+}
+
+#endif