generalized tangent vectors

gistfile1.txt

import TensorFlow

// MARK: - Protocols

protocol Vector {
  func scaled(by factor: Float) -> Self
  func adding(_ other: Self) -> Self
  static var zero: Self { get }
}

///
protocol TangentKind {
  associatedtype TFloat: Vector
  associatedtype TFloat2: Vector

  static func project0(_ t: TFloat2) -> TFloat
  static func project1(_ t: TFloat2) -> TFloat
  static func combine(_ t0: TFloat, _ t1: TFloat) -> TFloat2
}

// MARK: - Derivative functions

func product<K: TangentKind>(_ x: Float, _ y: Float, _ tIn: K.TFloat2, _ k: K.Type) -> (value: Float, tOut: K.TFloat) {
  (
    x * y,
    K.project0(tIn).scaled(by: y).adding(K.project1(tIn).scaled(by: x))
  )
}

func sum<K: TangentKind>(_ x: Float, _ y: Float, _ tIn: K.TFloat2, _ k: K.Type) -> (value: Float, tOut: K.TFloat) {
  (
    x + y,
    K.project0(tIn).adding(K.project1(tIn))
  )
}

func square<K: TangentKind>(_ x: Float, _ tIn: K.TFloat, _ k: K.Type) -> (value: Float, tOut: K.TFloat) {
  product(x, x, K.combine(tIn, tIn), k)
}

@_specialize(where K == DelayedReverse<ForwardTangent, ForwardTangent.TFloat>)
func example1<K: TangentKind>(_ x: Float, _ tIn: K.TFloat, _ k: K.Type) -> (value: Float, tOut: K.TFloat) {
  // x^2 + 3 * x

  // x^2
  let (v1, t1) = square(x, tIn, k)

  // 3 * x
  let (v2, t2) = product(3, x, K.combine(K.TFloat.zero, tIn), k)

  // x^2 + 3 * x
  return sum(v1, v2, K.combine(t1, t2), k)
}

// MARK: - Tangent kinds

enum ForwardTangent: TangentKind {
  struct TFloat: Vector {
    var t: Float
    init(_ t: Float) { self.t = t }
    func scaled(by factor: Float) -> Self { Self(factor * t) }
    func adding(_ other: Self) -> Self { Self(t + other.t) }
    static var zero: Self { TFloat(0) }
  }
  struct TFloat2: Vector {
    var t0: Float
    var t1: Float
    init(_ t0: Float, _ t1: Float) { self.t0 = t0; self.t1 = t1 }
    func scaled(by factor: Float) -> Self { Self(factor * t0, factor * t1) }
    func adding(_ other: Self) -> Self { Self(t0 + other.t0, t1 + other.t1) }
    static var zero: Self { TFloat2(0, 0) }
  }

  static func project0(_ t: TFloat2) -> TFloat { TFloat(t.t0) }
  static func project1(_ t: TFloat2) -> TFloat { TFloat(t.t1) }
  static func combine(_ t0: TFloat, _ t1: TFloat) -> TFloat2 { TFloat2(t0.t, t1.t) }
}

enum DelayedReverse<K: TangentKind, In: Vector>: TangentKind {
  struct Pullback<Out: Vector>: Vector {
    var f: (Out) -> In
    init(_ f: @escaping (Out) -> In) { self.f = f }
    func scaled(by factor: Float) -> Self {
      Self({ f($0.scaled(by: factor)) })
    }
    func adding(_ other: Self) -> Self {
      Self({ f($0).adding(other.f($0)) })
    }
    static var zero: Self { Self({ _ in In.zero }) }
  }

  typealias TFloat = Pullback<K.TFloat>
  typealias TFloat2 = Pullback<K.TFloat2>

  static func project0(_ t: TFloat2) -> TFloat {
    TFloat({ t.f(K.combine($0, .zero)) })
  }
  static func project1(_ t: TFloat2) -> TFloat {
    TFloat({ t.f(K.combine(.zero, $0)) })
  }
  static func combine(_ t0: TFloat, _ t1: TFloat) -> TFloat2 {
    TFloat2({ t0.f(K.project0($0)).adding(t1.f(K.project1($0))) })
  }
}

enum ForwardJacobian: TangentKind {
  struct TFloat: Vector {
    // Shape [1, n].
    var t: Tensor<Float>
    init(_ t: Tensor<Float>) { self.t = t }
    func scaled(by factor: Float) -> Self { Self(factor * t) }
    func adding(_ other: Self) -> Self { Self(t + other.t) }
    static var zero: Self { Self(Tensor(0)) }
  }
  struct TFloat2: Vector {
    // Shape [2, n]
    var t: Tensor<Float>
    init(_ t: Tensor<Float>) { self.t = t }
    func scaled(by factor: Float) -> Self {  Self(factor * t) }
    func adding(_ other: Self) -> Self { Self(t + other.t) }
    static var zero: Self { Self(Tensor(0)) }
  }

  static func project0(_ t: TFloat2) -> TFloat {
    TFloat(t.t[0])
  }
  static func project1(_ t: TFloat2) -> TFloat {
    TFloat(t.t[1])
  }
  static func combine(_ t0: TFloat, _ t1: TFloat) -> TFloat2 {
    TFloat2(Tensor(stacking: [t0.t, t1.t]))
  }
}

// MARK: - Example invocations

func example1ForwardDerivative(_ x: Float) -> Float {
  let (_, d) = example1(x, ForwardTangent.TFloat(1), ForwardTangent.self)
  return d.t
}

func example1ReverseDerivative(_ x: Float) -> Float {
  typealias Tan = DelayedReverse<ForwardTangent, ForwardTangent.TFloat>
  let (_, d) = example1(x, Tan.TFloat({ $0 }), Tan.self)
  return d.f(ForwardTangent.TFloat(1)).t
}

print(example1ForwardDerivative(2))
print(example1ReverseDerivative(2))

// TODOs:
// - can I have higher order differentiation?
// - measure how good it optimizes in the SIL
// - can I have a reverse mode jacobian?
// - is there a good way to add more types without explosion of associatedtypes
// - is there a way to add generic types?