public protocol Layer: Differentiable & KeyPathIterable
    where AllDifferentiableVariables: KeyPathIterable {
    /// The input type of the layer.
    associatedtype Input: Differentiable
    /// The output type of the layer.
    associatedtype Output: Differentiable

    /// Returns the output obtained from applying to an input.
    @differentiable(wrt: (self, input))
    func applied(to input: Input) -> Output
}

public extension Layer {
    func valueWithPullback(at input: Input)
        -> (output: Output,
            pullback: (Output.CotangentVector)
                -> (layerGradient: CotangentVector, inputGradient: Input.CotangentVector)) {
        let (out, pullback) = _valueWithPullback(at: self, input, in: Self.applied(to:))
        return (out, pullback)
    }
}

public struct Dense<Scalar>: Layer
    where Scalar : FloatingPoint & Differentiable & TensorFlowScalar {
    @differentiable(vjp: vjpWeight)
    public var weight: Tensor<Scalar>
    @differentiable(vjp: vjpBias)
    public var bias: Tensor<Scalar>
    public typealias Activation = @autodiff (Tensor<Scalar>) -> Tensor<Scalar>
    @noDerivative public let activation: Activation

    @usableFromInline func vjpWeight() -> (Tensor<Scalar>, (Tensor<Scalar>) -> CotangentVector) {
        return (weight, { v in CotangentVector(weight: v, bias: .zero) })
    }
    @usableFromInline func vjpBias() -> (Tensor<Scalar>, (Tensor<Scalar>) -> CotangentVector) {
        return (bias, { v in CotangentVector(weight: .zero, bias: v) })
    }

    // FIXME
    public var allKeyPaths: [PartialKeyPath<Dense>] {
        return [\Dense.weight, \Dense.bias]
    }

    @differentiable(wrt: (self, input), vjp: _vjpApplied(to:))
    public func applied(to input: Tensor<Scalar>) -> Tensor<Scalar> {
        return activation(matmul(input, weight) + bias)
    }

    @usableFromInline
    func _vjpApplied(to input: Tensor<Scalar>)
        -> (Tensor<Scalar>, (Tensor<Scalar>) -> (Dense.CotangentVector, Tensor<Scalar>)) {
      let r0 = matmul(input, weight)
      let r1 = r0 + bias
      let biasShape = bias.shapeTensor
      func pullback(_ v: Tensor<Scalar>) -> (Dense.CotangentVector, Tensor<Scalar>) {
          return (Dense.CotangentVector(
                      weight: matmul(input.transposed(), v),
                      bias: v.unbroadcast(toShape: biasShape)),
                  matmul(v, weight.transposed()))
      }
      return (r1, pullback)
    }
}

Unhandled use of FunctionRefInst
UNREACHABLE executed at /usr/local/src/swift-build/swift/lib/IRGen/LoadableByAddress.cpp:2743!
Stack dump:
0.  Program arguments: /Library/Developer/Toolchains/swift-tensorflow-DEVELOPMENT-2019-01-21-a.xctoolchain/usr/bin/swift -frontend -c -primary-file /Users/rxwei/Development/ML/DeepLearning/Sources/DeepLearning/Layer.swift /Users/rxwei/Development/ML/DeepLearning/Sources/DeepLearning/Loss.swift /Users/rxwei/Development/ML/DeepLearning/Sources/DeepLearning/Optimizer.swift /Users/rxwei/Development/ML/DeepLearning/Sources/DeepLearning/Utilities.swift -emit-module-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer~partial.swiftmodule -emit-module-doc-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer~partial.swiftdoc -serialize-diagnostics-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer.dia -emit-dependencies-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer.d -emit-reference-dependencies-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer.swiftdeps -target x86_64-apple-macosx10.10 -enable-objc-interop -sdk /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.14.sdk -I /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Products/Debug -F /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Products/Debug -F /Applications/Xcode.app/Contents/Developer/Platforms/MacOSX.platform/Developer/Library/Frameworks -enable-testing -g -module-cache-path /Users/rxwei/Library/Developer/Xcode/DerivedData/ModuleCache.noindex -swift-version 4.2 -enforce-exclusivity=checked -Onone -D SWIFT_PACKAGE -D DEBUG -D Xcode -serialize-debugging-options -Xcc -working-directory -Xcc /Users/rxwei/Development/ML/DeepLearning -Xcc -I/Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/swift-overrides.hmap -Xcc -I/Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Products/Debug/include -Xcc -I/Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/DerivedSources/x86_64 -Xcc -I/Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/DerivedSources -Xcc -DDEBUG=1 -module-name DeepLearning -o /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Build/Intermediates.noindex/DeepLearning.build/Debug/DeepLearning.build/Objects-normal/x86_64/Layer.o -index-store-path /Users/rxwei/Library/Developer/Xcode/DerivedData/DeepLearning-ccwrhnlcxxkamxedajvpoaummgzs/Index/DataStore -index-system-modules 
1.  While running pass #​71 SILModuleTransform "LoadableByAddress".
0  swift                    0x00000001136e80b8 llvm::sys::PrintStackTrace(llvm::raw_ostream&) + 40
1  swift                    0x00000001136e7335 llvm::sys::RunSignalHandlers() + 85
2  swift                    0x00000001136e86c2 SignalHandler(int) + 258
3  libsystem_platform.dylib 0x00007fff72317b3d _sigtramp + 29
4  swift                    0x0000000114537008 (anonymous namespace)::DarwinX86AsmBackend::getCompactUnwindRegNum(unsigned int) const::CU64BitRegs + 239278
5  libsystem_c.dylib        0x00007fff721d51c9 abort + 127
6  swift                    0x0000000113683024 llvm::llvm_unreachable_internal(char const*, char const*, unsigned int) + 532
7  swift                    0x000000010fbe2bf7 (anonymous namespace)::LoadableByAddress::run() + 13015
8  swift                    0x000000011064943b swift::SILPassManager::runModulePass(unsigned int) + 1083
9  swift                    0x0000000110649f14 swift::SILPassManager::execute() + 692
10 swift                    0x000000010fb926bb swift::SILPassManager::executePassPipelinePlan(swift::SILPassPipelinePlan const&) + 187
11 swift                    0x000000010fb92112 runIRGenPreparePasses(swift::SILModule&, swift::irgen::IRGenModule&) + 162
12 swift                    0x000000010fb9066e performIRGeneration(swift::IRGenOptions&, swift::ModuleDecl*, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::StringRef, swift::PrimarySpecificPaths const&, llvm::LLVMContext&, swift::SourceFile*, llvm::GlobalVariable**) + 894
13 swift                    0x000000010fb90c56 swift::performIRGeneration(swift::IRGenOptions&, swift::SourceFile&, std::__1::unique_ptr<swift::SILModule, std::__1::default_delete<swift::SILModule> >, llvm::StringRef, swift::PrimarySpecificPaths const&, llvm::LLVMContext&, llvm::GlobalVariable**) + 86
14 swift                    0x000000010fa1ad2b performCompile(swift::CompilerInstance&, swift::CompilerInvocation&, llvm::ArrayRef<char const*>, int&, swift::FrontendObserver*, swift::UnifiedStatsReporter*) + 13739
15 swift                    0x000000010fa166ed swift::performFrontend(llvm::ArrayRef<char const*>, char const*, void*, swift::FrontendObserver*) + 3021
16 swift                    0x000000010f9c847e main + 686
17 libdyld.dylib            0x00007fff7212ced9 start + 1