Skip to content

Commit 87e4d99

Browse files
committed
Add missed derivative of Tracked<T>.*=
1 parent 89dba2b commit 87e4d99

File tree

1 file changed

+23
-0
lines changed

1 file changed

+23
-0
lines changed

stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift.gyb

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,29 @@ where
340340
{
341341
return (lhs * rhs, { (dx, dy) in dx * rhs + dy * lhs })
342342
}
343+
344+
@usableFromInline
345+
@derivative(of: *=)
346+
static func _vjpMultiplyAssign(_ lhs: inout Self, _ rhs: Self) -> (
347+
value: Void, pullback: (inout Self) -> Self)
348+
{
349+
defer { lhs *= rhs }
350+
return ((), { [lhs = lhs] v in
351+
let drhs = lhs * v
352+
v *= rhs
353+
return drhs
354+
})
355+
}
356+
357+
@usableFromInline
358+
@derivative(of: *=)
359+
static func _jvpMultiplyAssign(_ lhs: inout Self, _ rhs: Self) -> (
360+
value: Void, differential: (inout Self, Self) -> Void)
361+
{
362+
let oldLhs = lhs
363+
lhs *= rhs
364+
return ((), { $0 = $0 * rhs + oldLhs * $1 })
365+
}
343366
}
344367

345368
extension ${Self}

0 commit comments

Comments
 (0)