Example wanted: manually defined tangent vectors

[BradLarson]

We should have an example of what's required to define a manual tangent vector, which can be tricky to set up right.

[pedronahum]

A proposal to close this issue:

import _Differentiation

// A 2D point that we want to differentiate through.
struct MyPoint: Differentiable {
    var x: Float
    var y: Float

    // Custom tangent type. Must have fields that match
    // the stored properties 'x' and 'y' of MyPoint.
    struct TangentVector: Differentiable & AdditiveArithmetic {
        var x: Float
        var y: Float

        // Required by AdditiveArithmetic.
        static var zero: TangentVector {
            TangentVector(x: 0, y: 0)
        }

        static func + (lhs: TangentVector, rhs: TangentVector) -> TangentVector {
            TangentVector(x: lhs.x + rhs.x, y: lhs.y + rhs.y)
        }

        static func - (lhs: TangentVector, rhs: TangentVector) -> TangentVector {
            TangentVector(x: lhs.x - rhs.x, y: lhs.y - rhs.y)
        }
    }

    // The Differentiable protocol expects 'move(by:)' to be defined.
    mutating func move(by offset: TangentVector) {
        x += offset.x
        y += offset.y
    }
}

// Marked differentiable with reverse-mode AD.
@differentiable(reverse)
func squaredDistanceFromOrigin(_ point: MyPoint) -> Float {
    point.x * point.x + point.y * point.y
}

func runExample() {
    let initialPoint = MyPoint(x: 3, y: 4)
    // Compute value and gradient at `initialPoint`.
    let (value, gradient) = valueWithGradient(
        at: initialPoint,
        of: squaredDistanceFromOrigin
    )

    print("Value at (\(initialPoint.x), \(initialPoint.y)): \(value)")
    // The gradient is now a TangentVector with matching field names.
    print("Gradient: (x: \(gradient.x), y: \(gradient.y))")

    // Move the point along the gradient direction.
    var movedPoint = initialPoint
    movedPoint.move(by: gradient)
    print("Moved point: (x: \(movedPoint.x), y: \(movedPoint.y))")
}

runExample()

}

And a motivating example that shows that if you rely on auto-synthesis of the tangent vector but your struct actually has fewer (or differently named) stored properties than you think, you might get a gradient that is mathematically correct for the underlying storage, but surprising if you conceptualized your type as having multiple independent coordinates.

import _Differentiation

/// A struct that *appears* to have two coordinates, x and y,
/// but actually has only one stored property, `storage`.
/// We do NOT define our own TangentVector here, so Swift auto-synthesizes one.
/// That auto-synth is based on the single stored property 'storage',
/// yielding a 1D tangent vector.
struct MyPoint: Differentiable {
    // Only one real stored property.
    private var storage: Float

    // x is computed directly from `storage`.
    var x: Float {
        get { storage }
        set { storage = newValue }
    }

    // y is also computed from the same `storage`, shifted by +10.
    var y: Float {
        get { storage + 10 }
        set { storage = newValue - 10 }
    }

    // Because we do NOT provide a custom TangentVector and do NOT override
    // `move(by:)`, Swift will auto-synthesize something like:
    //
    //   struct TangentVector: Differentiable, AdditiveArithmetic {
    //       var storage: Float
    //       ...
    //   }
    //
    // i.e., a *single* float storing the derivative of `storage`.
    //
    // Meanwhile, someone reading code that references x and y might
    // mistakenly think x and y are truly independent.

    init(x: Float, y _: Float) {
        // Contrive some setup. For example, always let storage = x,
        // ignoring the 'y' argument or combining them in some custom way.
        storage = x
    }
}

@differentiable(reverse)
func distanceFromOrigin(_ point: MyPoint) -> Float {
    // Looks like normal 2D distance, but under the hood, there's only one
    // stored property. Swift sees just one dimension's worth of "movement."
    point.x * point.x + point.y * point.y
}

func runAutoSynthExample() {
    // Initialize a MyPoint that suggests x=3, y=4.
    // But we ignore 'y' in the init? Let's say we do x=3, y=4 anyway:
    var p = MyPoint(x: 3, y: 4)
    // Let's see what p.x and p.y actually are:
    print("p.x = \(p.x), p.y = \(p.y)")

    // Compute the value and gradient at p.
    // If the code compiles, Swift auto-synthesizes a single-field tangent vector.
    let (value, grad) = valueWithGradient(at: p, of: distanceFromOrigin)

    print("Value of distance = \(value)")
    // The gradient is presumably a single struct with a single property,
    // e.g. "TangentVector(storage: ...)"
    // which merges the notion of changes to x and y.
    print("Gradient = \(grad)")

    // Attempt to move p by that gradient:
    p.move(by: grad)
    print("After moving by the gradient => p.x = \(p.x), p.y = \(p.y)")
}

runAutoSynthExample()