[SR-14113] Support _read and _modify accessor differentiation

[dan-zheng]

Previous ID	SR-14113
Radar	rdar://problem/73742489
Original Reporter	@dan-zheng
Type	Task

Additional Detail from JIRA

Votes	3
Component/s	Compiler
Labels	Task, AutoDiff
Assignee	None
Priority	Medium

md5: fc0ab68747451d6712b8e1289f2c8b68

Sub-Tasks:

• SR-14114 Diagnose coroutine differentiation until it is supported
• SR-14115 SIL: derivative type calculation for coroutines
• SR-14116 Differentiation: class property modify accessor applications are not active
• SR-14117 stdlib: support differentiation of Array.subscript._modify

blocks:

• SR-12640 Differentiation transform: support wrapped value modify accessors

is blocked by:

• TF-129 Support differentiation of functions with inout parameters

relates to:

• TF-1078 [AD] Incorrect derivatives for coroutines (begin_apply for a modify accessor)

Issue Description:

Support differentiation of coroutines. read and modify accessors are coroutines.
SIL has dedicated coroutine function types: https://github.com/apple/swift/blob/master/docs/SIL.rst#coroutine-types

Consider adding subtasks when starting work.

modify accessor differentiation is blocked by TF-129: inout argument differentiation. modify accessors have inout arguments.

[swift-ci]

Comment by Pedro N. Rodriguez (JIRA)

Hi @dan-zheng

In https://bugs.swift.org/browse/TF-1078, a workaround was recommended to address the issue. Using v0.7-rc2 Ubuntu CPU Only, the following code will fail to compile with an error suggesting that the function is not differentiable.

{code:Swift}

import TensorFlow

extension Array {
/// A functional version of `Array.subscript.modify`.
/// Differentiation does yet not support `Array.subscript.modify` because
/// it is a coroutine.
@differentiable(where Element: Differentiable)
func updated(at index: Int, with newValue: Element) -> [Element] {
var result = self
result[index] = newValue
return result
}
}

extension Array where Element: Differentiable {
/// An inefficient derivative for a functional version of `Array.subscript.modify`.
@Derivative(of: updated)
func vjpUpdated(at index: Int, with newValue: Element) -> (
value: [Element], pullback: (TangentVector) -> (TangentVector, Element.TangentVector)
) {
let value = updated(at: index, with: newValue)
return (value, { [count = count] v in
// `dself`: `v` but with `v[index] = 0`.
// `dnewValue`: `v[index]`.
var dself = v
dself.base[index] = .zero
return (dself, v[index])
})
}
}

struct PathCalc : Differentiable {
@noDerivative let N: Int
@noDerivative let Nmat: Int
@noDerivative let Delta: Float
@noDerivative let Z: Float

@differentiable(wrt: (Rates, Lambda))
func applied(Rates:Array<Float>,Lambda:Array<Float> ) -> Array<Float> {

var L = Array<Float>(repeating: 0.0,count: N)
let sq = Delta.squareRoot()

for n in 0...Nmat-1{
let sqez = sq*Z
var v = Float(0)

for i in n+1...N-1{
let lam = Lambda[i - n - 1 ]
let con1 = Delta*lam
v = v + (con1*Rates[i])/(1.0 + Delta*L[i])
let first = con1*v
let second = lam*(sqez - 0.5*con1)
let vrat = exp(first + second)
L[i] = Rates[i]*vrat

}
}
return L
}

}

struct Portfolio: Differentiable {

@noDerivative let N: Int
@noDerivative let Nmat: Int
@noDerivative let Nopt: Int
@noDerivative let Delta: Float
@noDerivative let Maturities: Array<Int>
@noDerivative let SwapRates: Array<Float>
@noDerivative let Z: Float

@differentiable(wrt: (self, Rates, Lambda))
func calculate(Rates:Array<Float>,Lambda:Array<Float>) -> Float {

let pathCalcModel = PathCalc(N: N, Nmat: Nmat, Delta: Delta, Z:Z)
let L = pathCalcModel.applied(Rates:Rates , Lambda:Lambda )

var S = Array<Float>(repeating: 0.0,count: N)
var B = Array<Float>(repeating: 0.0,count: N)

var b = Float(1.0)
var s = Float(0)

for n in Nmat...N-1{
b = b / (1.0 + Delta*L[n])
s = s + Delta*b
B[n] = b
S[n] = s
}

var v = Float(0)
var swapval = Float(0)

for i in 0...Nopt-1{
let m = Int(Maturities[i] + Nmat - 1)
swapval = B[m] + SwapRates[i]*S[m] - 1.0
if (swapval < 0.0){
v = v + -100.0*swapval
}
}
// Apply discount
for n in 0...Nmat-1{
v = v / (1.0 + Delta*L[n])

}

return v
}

}

let Nmat = 40
let N = Nmat + 40
let Nopt = 15
let maturities = Array<Int>([4,4,4,8,8,8,20,20,20,28,28,28,40,40,40])
let swapRates = Array<Float>([0.045,0.05,0.055,0.045,0.05,0.055,0.045,0.05,
0.055,0.045,0.05,0.055,0.045,0.05,0.055])

let rates = Array<Float>(repeating: 0.05,count: N)
let lambdas = Array<Float>(repeating: 0.2,count: N)

let portfolioModel = Portfolio(N:N , Nmat:Nmat , Nopt: Nopt , Delta:0.25,
Maturities: maturities , SwapRates:swapRates , Z:0.3 )
let portfolioValue = portfolioModel.calculate(Rates:rates , Lambda:lambdas )
print(portfolioValue)

let (g_model, g_rates,g_lambdas) = gradient(at: portfolioModel, rates, lambdas)
{ portfolioModel, rates, lambdas in
portfolioModel.calculate(Rates: rates, Lambda:lambdas)
}
print(g_model)
print(g_rates)

{code}

[porterchild]

You have to use the workaround's .updated() function explicitly, it does not override subscripting so that things will work automatically. Here is an updated version of two of your for loops that need to use the workaround:

for i in n+1...N-1{
    let lam = Lambda[i - n - 1 ]
    let con1 = Delta*lam
    v = v + (con1*Rates[i])/(1.0 + Delta*L[i])
    let first = con1*v
    let second = lam*(sqez - 0.5*con1)
    let vrat = exp(first + second)
    //L[i] = Rates[i]*vrat
    L = L.updated(at: i, with: Rates[i] * vrat)
}

and

for n in Nmat...N-1{
    b = b / (1.0 + Delta*L[n])
    s = s + Delta*b
    //B[n] = b
    B = B.updated(at: n, with: b)
    //S[n] = s
    S = S.updated(at: n, with: s)
}

[porterchild]

Also I needed to add an update to the workaround, I'll go comment on TF-1078 with that change.

[swift-ci]

Comment by Pedro N. Rodriguez (JIRA)

@porterchild, many thanks for reviewing the code After implementing the suggested changes, the code runs smoothly.

[dabrahams]

IIUC, _read isn't a thing; it's just left over or something. I suppose I could be reading into this post: https://forums.swift.org/t/what-does-collection-subscript-read-do-in-swift-5/20014/3

[rxwei]

@swift-ci create

[rxwei]

@swift-ci create

[philipturner]

@dan-zheng I do not think it's possible to find bug TF-129, given that the subproject died and (maybe) some bugs transferred over the Apple Sync System. Since you authored this, could you explain that bug in greater detail or show an SR bug it maps to?

[dan-zheng]

@philipturner According to the description of this issue, TF-129 tracked supporting differentiation of inout arguments.

I couldn't find a currently existing bug corresponding to TF-129, but inout argument differentiation is supported (implementation) and doesn't block this issue.