Simplify API for getting layers by key

Author: shadaj

Examples/LeNet-MNIST/main.swift

@@ -41,7 +41,7 @@ let denseDef = AutoConv2D<Float>(filterShape: (5, 5), outputChannels: 6, padding
   .then(AutoDense(outputSize: 84, activation: relu))
   .then(KeyedAutoLayer(AutoDense(outputSize: 10), key: lastDenseKey))
 
-var (classifier, keys) = denseDef.buildModelWithKeys(inputShape: (28, 28, 1))
+var classifier = denseDef.buildModel(inputShape: (28, 28, 1))
 classifier.move(to: device)
 
 var optimizer = SGD(for: classifier, learningRate: 0.1)
@@ -123,5 +123,5 @@ for (epoch, epochBatches) in dataset.training.prefix(epochCount).enumerated() {
         (\(String(format: "%.3f", testStats.accuracy))%)
         """)
 
-    print(lastDenseKey.readFrom(layerInstance: classifier, keyDict: keys).bias)
+    print(classifier[lastDenseKey].bias)
 }

Models/ImageClassification/ResNet.swift

@@ -125,7 +125,7 @@ public func autoResNet(
 
 /// An implementation of the ResNet v1 and v1.5 architectures, at various depths.
 public struct ResNet: Layer {
-    public var underlying: AutoResNet.InstanceType
+    public var underlying: BuiltAutoLayer<AutoResNet.InstanceType>
 
     /// Initializes a new ResNet v1 or v1.5 network model.
     ///

Models/LayerInit/AutoLayer.swift

@@ -9,10 +9,10 @@ public protocol AutoLayer {
 }
 
 extension AutoLayer {
-    public func buildModel(inputShape: InputShape) -> InstanceType {
+    public func buildModel(inputShape: InputShape) -> BuiltAutoLayer<InstanceType> {
         var keyDict: [AnyAutoLayerKey: Any] = [:]
         let (layerInstance, _) = self.buildModelWithOutputShape(inputShape: inputShape, keyPathSoFar: \InstanceType.self, keyDict: &keyDict)
-        return layerInstance
+        return BuiltAutoLayer(layer: layerInstance, keyMapping: keyDict)
     }
 
     public func buildModelWithKeys(inputShape: InputShape) -> (InstanceType, [AnyAutoLayerKey: Any]) {
@@ -21,3 +21,22 @@ extension AutoLayer {
         return (layerInstance, keyDict)
     }
 }
+
+public struct BuiltAutoLayer<InstanceType: Layer>: Layer {
+    public var layer: InstanceType
+    @noDerivative let keyMapping: [AnyAutoLayerKey: Any]
+
+    public init(layer: InstanceType, keyMapping: [AnyAutoLayerKey: Any]) {
+        self.layer = layer
+        self.keyMapping = keyMapping
+    }
+
+    @differentiable
+    public func callAsFunction(_ input: InstanceType.Input) -> InstanceType.Output {
+        return layer(input)
+    }
+
+    public subscript<T>(index: AutoLayerKey<T>) -> T {
+        return self.layer[keyPath: self.keyMapping[index] as! KeyPath<InstanceType, T>]
+    }
+}

Models/LayerInit/AutoLayerKey.swift

@@ -14,12 +14,6 @@ public class AutoLayerKey<T: Layer>: AnyAutoLayerKey {
     public override init() {}
 }
 
-extension AutoLayerKey {
-    public func readFrom<Instance: Layer>(layerInstance: Instance, keyDict: [AnyAutoLayerKey: Any]) -> T {
-        return layerInstance[keyPath: keyDict[self] as! KeyPath<Instance, T>]
-    }
-}
-
 public struct KeyedAutoLayer<Underlying: AutoLayer>: AutoLayer {
     let underlying: Underlying
     let key: AutoLayerKey<InstanceType>

Models/LayerInit/AutoModule.swift

@@ -6,7 +6,7 @@ public protocol AutoModule: AutoLayer {
     var initializeLayer: LayerType { mutating get }
 }
 
-extension AutoModule  {
+extension AutoModule {
     public typealias InstanceType = LayerType.InstanceType
     public typealias InputShape = LayerType.InputShape
     public typealias OutputShape = LayerType.OutputShape