Staged Y combinator with NNKit

staged-y-combinator.swift

/// Y combinator
func fix<D, R>(
    _ f: @escaping (Rep<(D) -> R>) -> Rep<(D) -> R>) -> Rep<(D) -> R> {
    return lambda { d in f(fix(f))[d] }
}

let fac: Rep<(Int) -> Int> = fix { f in
    lambda { (n: Rep<Int>) in
        .if(n == 0, then: ^1, else: n * f[n - 1])
    }
}