Swift: Differentiable KeyPath-based reads

[porterchild]

I'm trying to make object member accesses using KeyPaths differentiable, I think I'm pretty close on a workaround after adapting Marc's workaround for Array subscript reads.

I'm stuck where, inside my pullback for read(at:), I attempt a v[keyPath: key] = v.tempMember but the keyPath is a WritableKeyPath<ASDF, Float>, not a WritableKeyPath<ASDF.TangentVector, Float>. This causes a compiler error.

I'm wondering if there is a way around this, maybe through some EuclideanDifferentiable functionality? Or maybe KeyPathIterable?

Thanks!

// goal: make reads on an object through a keyPath differentiable
// ex: let a = object[keyPath: key]


//---------------------------------------------------------
struct ASDF: Differentiable{
    var a: Float
    var b: Float

    var tempMember: Float = 0

    @differentiable
    mutating func read(at key: WritableKeyPath<ASDF, Float>){
        self.tempMember = self[keyPath: key]
    }
}
extension ASDF{
    @derivative(of: read)
    mutating func vjpRead(at key: WritableKeyPath<ASDF, Float>) -> (value: Void, pullback: (inout ASDF.TangentVector) -> ()) 
    {
        func pullback(_ v: inout ASDF.TangentVector) {
            v[keyPath: key] = v.tempMember   //error: Key path with root type 'ASDF' cannot be applied to a base of type 'ASDF.TangentVector'
        }
        read(at: key)
        return ((), pullback)
    }
}



//---------------------------------------------------------
//usage

let aKey: WritableKeyPath<ASDF, Float> = \ASDF.a
let bKey: WritableKeyPath<ASDF, Float> = \ASDF.b

@differentiable
func multiplyMembers(asdf: ASDF) -> Float{
    var asdf = asdf

    asdf.read(at: aKey)
    let a = asdf.tempMember
    asdf.read(at: bKey)
    let b = asdf.tempMember
    let c = a * b

    //aka
//    let c = asdf[keyPath: aKey] * asdf[keyPath: bKey]

    return c
}

print(multiplyMembers(asdf: ASDF(a: 2, b: 3)))
print(gradient(of: multiplyMembers)(ASDF(a: 2, b: 3)))

[dan-zheng]

Key paths have a statically fixed Root type, so you can apply a PartialKeyPath<T> to a value of type T.TangentVector iff T == T.TangentVector.

---

Quick fix for your example: you can declare a ASDF: AdditiveArithmetic conformance.

The Differentiable.TangentVector associated type is required to conform to AdditiveArithmetic. Declaring a ASDF: AdditiveArithmetic conformance causes the compiler to infer ASDF.TangentVector == ASDF during Differentiable conformance derivation.

With this, WritableKeyPath<ASDF, Float> can be applied to both ASDF and ASDF.TangentVector values.

struct ASDF: Differentiable & AdditiveArithmetic {
  var a: Float
  var b: Float

  var tempMember: Float = 0

  @differentiable
  mutating func read(at key: WritableKeyPath<ASDF, Float>) {
    self.tempMember = self[keyPath: key]
  }
}

extension ASDF {
  @derivative(of: read)
  mutating func vjpRead(at key: WritableKeyPath<ASDF, Float>) -> (
    value: Void, pullback: (inout ASDF.TangentVector) -> Void
  ) {
    func pullback(_ v: inout ASDF.TangentVector) {
      v[keyPath: key] = v.tempMember
    }
    read(at: key)
    return ((), pullback)
  }
}

@differentiable
func multiplty(_ asdf: ASDF) -> Float {
  var asdf = asdf

  asdf.read(at: \ASDF.a)
  let a = asdf.tempMember
  asdf.read(at: \ASDF.b)
  let b = asdf.tempMember

  let c = a * b
  return c
}

print(valueWithGradient(of: multiplty)(ASDF(a: 2, b: 3)))
// (value: 6.0, gradient: kp.ASDF(a: 5.0, b: 2.0, tempMember: 5.0))

[dan-zheng]

Here's a more generic version:

import _Differentiation

extension Differentiable {
  @differentiable(where Self == TangentVector, T: Differentiable, T == T.TangentVector)
  mutating func write<T>(
    from source: WritableKeyPath<Self, T>,
    to destination: WritableKeyPath<Self, T>
  ) {
    self[keyPath: destination] = self[keyPath: source]
  }

  @derivative(of: write)
  mutating func vjpWrite<T: Differentiable>(
    from source: WritableKeyPath<Self, T>,
    to destination: WritableKeyPath<Self, T>
  ) -> (
    value: Void, pullback: (inout Self) -> Void
  ) where Self == TangentVector, T == T.TangentVector {
    func pullback(_ dself: inout Self) {
      dself[keyPath: source] = dself[keyPath: destination]
    }
    self[keyPath: destination] = self[keyPath: source]
    return ((), pullback)
  }
}

struct Vector: Differentiable & AdditiveArithmetic {
  var x: Float
  var y: Float

  var tmp: Float = 0

  @differentiable
  mutating func read(at key: WritableKeyPath<Vector, Float>) {
    write(from: key, to: \Vector.tmp)
  }
}

func multiply(_ v: Vector) -> Float {
  var v = v
  v.read(at: \Vector.x)
  let x = v.tmp
  v.read(at: \Vector.y)
  let y = v.tmp
  return x * y
}
print(valueWithGradient(at: Vector(x: 2, y: 3), in: multiply))
// (value: 6.0, gradient: Vector(x: 5.0, y: 2.0, tmp: 5.0))

[porterchild]

Great! This genericized improvement lets me require much less from types that want to have the differentiable read() function. Thanks for the help Dan!

For anyone who finds this later, here is the complete version with reads and writes:
(and a read derivative accuracy update)

extension Differentiable {
    //-----------------------------------------------------------
    // internal write, if a type wants read() it has to make one and call this inside
    @differentiable(where Self == TangentVector, T: Differentiable, T == T.TangentVector)
    mutating func internalWrite<T>(from source: WritableKeyPath<Self, T>, to destination: WritableKeyPath<Self, T>) {
        self[keyPath: destination] = self[keyPath: source]
    }

    @derivative(of: internalWrite)
    mutating func vjpInternalWrite<T: Differentiable>(from source: WritableKeyPath<Self, T>, to destination: WritableKeyPath<Self, T>) ->
    (value: Void, pullback: (inout Self.TangentVector) -> Void) where Self == TangentVector, T == T.TangentVector
    {
        func pullback(_ dself: inout Self.TangentVector) {
            dself[keyPath: source] += dself[keyPath: destination]
            dself[keyPath: destination] = .zero
        }
        self[keyPath: destination] = self[keyPath: source]
        return ((), pullback)
    }

    //-----------------------------------------------------------
    // write
    @differentiable(where Self == TangentVector, T: Differentiable, T == T.TangentVector)
    mutating func write<T: Differentiable>(to member: WritableKeyPath<Self, T>, with value: T){
        self[keyPath: member] = value
    }

    @derivative(of: write)
    mutating func vjpWrite<T: Differentiable>(to member: WritableKeyPath<Self, T>, with value: T) ->
    (value: Void, pullback: (inout Self.TangentVector) -> T.TangentVector) where Self == TangentVector, T == T.TangentVector
    {
        func pullback(_ dself: inout Self.TangentVector) -> T.TangentVector{
            let dWriteValue = dself[keyPath: member]
            dself[keyPath: member] = .zero
            return dWriteValue
        }
        write(to: member, with: value)
        return ((), pullback)
    }

}



struct ASDF: Differentiable, AdditiveArithmetic{
    var a: Float
    var b: Float

    var tmp: Float = 0

    @differentiable
    mutating func read(at key: WritableKeyPath<Self, Float>) {
      internalWrite(from: key, to: \Self.tmp)
    }
}




//---------------------------------------------------------
//usage
let aKey: WritableKeyPath<ASDF, Float> = \ASDF.a
let bKey: WritableKeyPath<ASDF, Float> = \ASDF.b

@differentiable
func aTimesA(asdf: ASDF) -> Float{
    var asdf = asdf

    //let a = asdf[keyPath: aKey]
    asdf.read(at: aKey)
    let a = asdf.tmp

    //asdf[keyPath: bKey] = a
    asdf.write(to: bKey, with: a)

    //let b = asdf[keyPath: bKey]
    asdf.read(at: bKey)
    let b = asdf.tmp

    let c = a * b

    //aka
//    let a = asdf[keyPath: aKey]
//    asdf[keyPath: bKey] = a
//    let b = asdf[keyPath: bKey]
//    let c = a * b

    return c
}

//print("forward", aTimesA(asdf: ASDF(a: 2, b: 3)))

//this function is equivalent to a^2, so derivative wrt a should be 2a
print("val and gradient:", valueWithGradient(of: aTimesA)(ASDF(a: 5, b: 6)))

[porterchild]

I got so focused on the above approach (which has O(1) speed on the reverse pass) that I forgot there's a simpler version that lets you forget about the .tmp workaround, in exchange for O(n) speed on the reverse pass (which will eventually be O(1)). Here it is:

extension Differentiable {
    //-----------------------------------------------------------
    // a read that's free of the .tmp workaround, but O(n) on backwards pass (because of zeroTangentVector materialization)
    @inlinable
    @differentiable(where Self == TangentVector, T: Differentiable, T == T.TangentVector)
    public func readSlow<T: Differentiable>(at member: WritableKeyPath<Self, T>) -> T{
        return self[keyPath: member]
    }

    @inlinable
    @derivative(of: readSlow)
    public func vjpReadSlow<T: Differentiable>(at member: WritableKeyPath<Self, T>) -> (value: T, pullback: (T.TangentVector) -> Self.TangentVector)
    where Self == TangentVector, T == T.TangentVector
    {
        return (value: self[keyPath: member], pullback:{ downstream in
            var zeroes = self.zeroTangentVector
            zeroes[keyPath: member] = downstream
            return zeroes
        })
    }
}

Usage:

//instead of this:
asdf.read(at: bKey)
let b = asdf.tmp

//you can do this
let b = asdf.readSlow(at: bKey)