From ff2a6db017b40265f6e97038b2212e1fd00e787c Mon Sep 17 00:00:00 2001 From: Dave Abrahams Date: Fri, 11 Dec 2020 21:31:29 -0800 Subject: [PATCH 1/3] Add T->T versions of mapBuffers and compactMapBuffers these can preserve the ElementAPI parameter. --- .../SwiftFusion/Core/TypeKeyedArrayBuffers.swift | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift b/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift index 5d0a4a9a..15a1e352 100644 --- a/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift +++ b/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift @@ -87,6 +87,21 @@ extension TypeKeyedArrayBuffers { return _storage.allSatisfy { kv in other._storage[kv.key]?.count == kv.value.count } } + /// Returns a mapping from each key `k` of `self` into `bufferTransform(self[k])`. + public func mapBuffers( + _ bufferTransform: (AnyArrayBuffer) throws -> AnyArrayBuffer + ) rethrows -> TypeKeyedArrayBuffers { + try .init(_storage: _storage.mapValues(bufferTransform)) + } + + /// Returns a mapping from each key `k` of `self` into the corresponding array, transformed by + /// `bufferTransform`. + public func compactMapBuffers( + _ bufferTransform: (AnyArrayBuffer) throws -> AnyArrayBuffer? + ) rethrows -> TypeKeyedArrayBuffers { + try .init(_storage: _storage.compactMapValues(bufferTransform)) + } + /// Returns a mapping from each key `k` of `self` into `bufferTransform(self[k])`. public func mapBuffers( _ bufferTransform: (AnyArrayBuffer) throws -> AnyArrayBuffer From 348c9e77fe30076150432db87698658415f08340 Mon Sep 17 00:00:00 2001 From: Dave Abrahams Date: Fri, 11 Dec 2020 21:44:49 -0800 Subject: [PATCH 2/3] Begin TypeKeyedArrayBuffers+Vector --- .../TypeKeyedArrayBuffers+Vector.swift | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift diff --git a/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift b/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift new file mode 100644 index 00000000..dafa0e41 --- /dev/null +++ b/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift @@ -0,0 +1,64 @@ +// Copyright 2020 The SwiftFusion Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import PenguinStructures + +extension TypeKeyedArrayBuffers where ElementAPI == VectorArrayDispatch { + /// Returns the squared norm of `self`, where `self` is viewed as a vector in the vector space + /// direct sum of all the variables. + /// + /// Precondition: All the variables in `self` are vectors. + public var squaredNorm: Double { + anyBuffers.lazy.map { $0.dot($0) }.reduce(0, +) + } + + /// Returns the scalar product of `lhs` with `rhs`, where `rhs` is viewed as a vector in the + /// vector space direct sum of all the variables. + /// + /// Precondition: All the variables in `rhs` are vectors. + public static func * (_ lhs: Double, _ rhs: Self) -> Self { + let r = rhs.mapBuffers { lhs * $0 } + return r + } + + /// Returns the vector sum of `lhs` with `rhs`, where `lhs` and `rhs` are viewed as vectors in the + /// vector space direct sum of all the variables. + /// + /// Precondition: All the elements in `lhs` and `rhs` are vectors. `lhs` and `rhs` have + /// assignments for exactly the same sets of variables. + public static func + (_ lhs: Self, _ rhs: Self) -> Self { + lhs.updatedBuffers(homomorphicArgument: rhs, +) + } + + /* + /// Stores an `ErrorVector` for a factor of type `F`. + public mutating func store(_ value: F.ErrorVector, factorType _: F.Type) { + _ = _storage[ + Type.id, + // Note: This is a safe upcast. + default: AnyVectorArrayBuffer(ArrayBuffer()) + ].unsafelyAppend(value) + } + + /// Returns the `ErrorVector` from the `perFactorID`-th factor of type `F`. + public subscript(_ perFactorID: Int, factorType _: F.Type) + -> F.ErrorVector + { + let array = _storage[Type.id].unsafelyUnwrapped + return ArrayBuffer(unsafelyDowncasting: array).withUnsafeBufferPointer { b in + b[perFactorID] + } + } + */ +} From 3e2293103339de231c74124211e1aaddac44da31 Mon Sep 17 00:00:00 2001 From: Dave Abrahams Date: Sat, 12 Dec 2020 13:12:11 -0800 Subject: [PATCH 3/3] TypeKeyedArrayBuffers: Differentiable compiles --- .../Core/TypeKeyedArrayBuffers.swift | 21 ++++++ .../TypeKeyedArrayBuffers+Vector.swift | 69 +++++++++---------- 2 files changed, 52 insertions(+), 38 deletions(-) diff --git a/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift b/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift index 15a1e352..a65be30f 100644 --- a/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift +++ b/Sources/SwiftFusion/Core/TypeKeyedArrayBuffers.swift @@ -117,6 +117,27 @@ extension TypeKeyedArrayBuffers { try .init(_storage: _storage.compactMapValues(bufferTransform)) } + /// Returns the first key in `self` such that `predicate(self[k], other[k]) == true`, or `nil` if + /// no such key exists. + /// + /// - Requires: `self.hasSameStructure(as: parameter)` + public func firstBufferKey( + homomorphicArgument other: TypeKeyedArrayBuffers, + where predicate: ( + _ myBuffer: AnyArrayBuffer, + _ otherBuffer: AnyArrayBuffer) throws -> Bool + ) rethrows -> TypeID? { + precondition( + _storage.count == other._storage.count, + "parameter must have same structure as `self`") + return try _storage.first { kv in + guard let v1 = other._storage[kv.key] else { + fatalError("parameter must have same structure as `self`") + } + return try predicate(kv.value, v1) + }.map { $0.key } + } + /// Invokes `update` on each buffer of self, passing the buffer having the same `key` in /// `parameter` as a second argument. /// diff --git a/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift b/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift index dafa0e41..4528dd2a 100644 --- a/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift +++ b/Sources/SwiftFusion/Inference/TypeKeyedArrayBuffers+Vector.swift @@ -12,53 +12,46 @@ // See the License for the specific language governing permissions and // limitations under the License. +import _Differentiation import PenguinStructures -extension TypeKeyedArrayBuffers where ElementAPI == VectorArrayDispatch { - /// Returns the squared norm of `self`, where `self` is viewed as a vector in the vector space - /// direct sum of all the variables. - /// - /// Precondition: All the variables in `self` are vectors. - public var squaredNorm: Double { - anyBuffers.lazy.map { $0.dot($0) }.reduce(0, +) +extension TypeKeyedArrayBuffers: Equatable where ElementAPI: VectorArrayDispatch { + public static func == (lhs: Self, rhs: Self) -> Bool { + lhs.firstBufferKey(homomorphicArgument: rhs) { $0 != $1 } == nil } - - /// Returns the scalar product of `lhs` with `rhs`, where `rhs` is viewed as a vector in the +} + +extension TypeKeyedArrayBuffers: AdditiveArithmetic where ElementAPI == VectorArrayDispatch { + /// Returns the vector sum of `lhs` with `rhs`, where `lhs` and `rhs` are viewed as vectors in the /// vector space direct sum of all the variables. /// - /// Precondition: All the variables in `rhs` are vectors. - public static func * (_ lhs: Double, _ rhs: Self) -> Self { - let r = rhs.mapBuffers { lhs * $0 } - return r + /// Precondition: `lhs` and `rhs` have assignments for exactly the same sets of variables. + public static func + (lhs: Self, rhs: Self) -> Self { + lhs.updatedBuffers(homomorphicArgument: rhs) { $0 + $1 } } - /// Returns the vector sum of `lhs` with `rhs`, where `lhs` and `rhs` are viewed as vectors in the - /// vector space direct sum of all the variables. - /// - /// Precondition: All the elements in `lhs` and `rhs` are vectors. `lhs` and `rhs` have - /// assignments for exactly the same sets of variables. - public static func + (_ lhs: Self, _ rhs: Self) -> Self { - lhs.updatedBuffers(homomorphicArgument: rhs, +) + public static func += (lhs: inout Self, rhs: Self) { + lhs.updateBuffers(homomorphicArgument: rhs) { $0 += $1 } } - /* - /// Stores an `ErrorVector` for a factor of type `F`. - public mutating func store(_ value: F.ErrorVector, factorType _: F.Type) { - _ = _storage[ - Type.id, - // Note: This is a safe upcast. - default: AnyVectorArrayBuffer(ArrayBuffer()) - ].unsafelyAppend(value) + public static func - (lhs: Self, rhs: Self) -> Self { + lhs.updatedBuffers(homomorphicArgument: rhs) { $0 - $1 } } - - /// Returns the `ErrorVector` from the `perFactorID`-th factor of type `F`. - public subscript(_ perFactorID: Int, factorType _: F.Type) - -> F.ErrorVector - { - let array = _storage[Type.id].unsafelyUnwrapped - return ArrayBuffer(unsafelyDowncasting: array).withUnsafeBufferPointer { b in - b[perFactorID] - } + + public static func -= (lhs: inout Self, rhs: Self) { + lhs.updateBuffers(homomorphicArgument: rhs) { $0 -= $1 } + } + + public static var zero: Self { .init() } +} + +extension TypeKeyedArrayBuffers: Differentiable where ElementAPI: DifferentiableArrayDispatch { + public typealias TangentVector = MappedArrayBuffers + public mutating func move(along offset: TangentVector) { + updateBuffers(homomorphicArgument: offset) { $0.move(along: $1) } + } + + public var zeroTangentVectorInitializer: () -> TangentVector { + { .init() } } - */ }