[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 #54403
• [SR-14115] SIL: derivative type calculation for coroutines #54404
• [SR-14116] Differentiation: class property modify accessor applications are not active #54685
• [SR-14117] stdlib: support differentiation of Array.subscript._modify #55256

blocks:

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

is blocked by:

• [AutoDiff] Support inout argument differentiation. #30013

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.

[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.

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)

[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.

[asl]

From the practical standpoint, I think in general it makes sense to support only coroutines that are used as _read and _modify accessors. Then probably Array.subscript._modify and Dictionary.subscript._modify are those examples that we'd need to support and this would probably cover 95% of cases. Here are their implementations (https://github.com/apple/swift/blob/main/stdlib/public/core/Array.swift#L751 and https://github.com/apple/swift/blob/main/stdlib/public/core/Dictionary.swift#L897 - one of):

    _modify {
      _makeMutableAndUnique() // makes the array native, too
      _checkSubscript_mutating(index)
      let address = _buffer.mutableFirstElementAddress + index
      defer { _endMutation() }
      yield &address.pointee
    }

    _modify {
      let (bucket, found) = _variant.mutatingFind(key)
      let native = _variant.asNative
      if !found {
        let value = defaultValue()
        native._insert(at: bucket, key: key, value: value)
      }
      let address = native._values + bucket.offset
      defer { _fixLifetime(self) }
      yield &address.pointee
    }

They all look pretty same and therefore I think it does make sense to make the following restrictions of the coroutines, their types and their callers:

1. They should be @yield_once coroutines
2. We will not support abort_apply instruction (so, unwind part of coroutine body could be ignored)
3. For the pullback generation we only support coroutines with single yield instruction, disallowing things like

if (foo) {
  yield &something
} else
  yield &other
}

I think first two conditions are always true for _modify / _read accessors. The third one could easily be enforced as we will need to "split" the function into two parts: before yield and resume.

Having these restrictions we can semantically "linearize" control flow of co-routine and for the purposes of differentiation consider it as a function with multiple semantic results. In particular, _modify accessor from #54404 could have an (@inout Array<τ_0_0>.TangentVector, @inout_aliasable τ_0_0.) -> Void pullback type. _read accessors seem to be easier as it seems we can ignore resume part of the function – it is never active.

There might be some rough edges here and there. One of the things that I could think of is that we can yield the local value, e.g.:

_modify {
  var y : Float;
  yield &y;
}

Maybe we can diagnose such cases as either the y will be not active, or otherwise there might be some active stuff in the resume part.

Tagging @BradLarson