Skip to content

Commit a5ef5fe

Browse files
authored
Improve differentiableMap vjp performance (#83807)
Improve differentiableMap vjp performance by reserving capacity for the values and pullbacks arrays.
1 parent d837d60 commit a5ef5fe

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

stdlib/public/Differentiation/ArrayDifferentiation.swift

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -369,8 +369,11 @@ extension Array where Element: Differentiable {
369369
value: [Result],
370370
pullback: (Array<Result>.TangentVector) -> Array.TangentVector
371371
) {
372+
let count = self.count
372373
var values: [Result] = []
373374
var pullbacks: [(Result.TangentVector) -> Element.TangentVector] = []
375+
values.reserveCapacity(count)
376+
pullbacks.reserveCapacity(count)
374377
for x in self {
375378
let (y, pb) = valueWithPullback(at: x, of: body)
376379
values.append(y)

0 commit comments

Comments
 (0)