swiftswift5autodiff

How is the pullback defined for a function with an optional argument?


Here is a minimal example:

import _Differentiation

@differentiable(reverse)
func g(x: Double?) -> Double {
    if x == nil {
        return 0.0  // I don't care what this value is
    }
    else {
        return x! * x!
    }
}

@derivative(of: g)
func gVJP(x: Double?) -> (value: Double, pullback: (Double) -> Double?) {
    let value = g(x: x)
    
    func pullback(_ dOutput: Double) -> Double? {
        if x == nil {
            return nil // I don't care what this value is
        }
        else {
            return dOutput * 2.0 * x!
        }
    }
    
    return (value: value, pullback: pullback)
}

I get the following compiler error:

Function result's 'pullback' type does not match 'g(x:)'
1. 'pullback' does not have expected type '(Double.TangentVector) -> Optional<Double>.TangentVector' (aka '(Double) -> Optional<Double>.TangentVector')

I tried defining the return type of the pullback to be Optional<Double>.TangentVector, but that gave me errors that my pullback's return values were incompatible with return type 'Optional<Double>.TangentVector'.


Solution

  • The return values of pullback need to be cast to type Optional<Double>.TangentVector:

    import _Differentiation
    
    @differentiable(reverse)
    func g(x: Double?) -> Double {
        if x == nil {
            return 0.0  // I don't care what this value is
        }
        else {
            return x! * x!
        }
    }
    
    @derivative(of: g)
    func gVJP(x: Double?) -> (value: Double, pullback: (Double) -> Optional<Double>.TangentVector) {
        let value = g(x: x)
    
        func pullback(_ dOutput: Double) ->Optional<Double>.TangentVector {
            if x == nil {
                return Optional<Double>.TangentVector(nil) // I don't care what this value is
            }
            else {
                return Optional<Double>.TangentVector(dOutput * 2.0 * x!)
            }
        }
    
        return (value: value, pullback: pullback)
    }