Simplify API for getting layers by key

Author: shadaj

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