Simplify ResNet block definition

Author: shadaj

Models/ImageClassification/ResNet.swift

@@ -66,17 +66,17 @@ public func autoResidualBlock(inputFilters: Int, filters: Int, strides: (Int, In
     }
 
     let earlyConvsWithRelu = earlyConvs.map({ (conv) in
-        conv.then(AutoFunction(fnShape: { $0 }, fn: { (prev: Tensor<Float>) in relu(prev) }))
+        conv.then(AutoFunction(fnShape: { $0 }, fn: { relu($0) }))
     })
 
     let lastConvResult = AutoSequencedMany(layers: earlyConvsWithRelu).then(lastConv)
 
     let convPlusResidual = AutoSplitMerge(
         layer1: residual,
         layer2: lastConvResult,
-        mergeOutputShape: { (l1, l2) in l1 }, mergeFn: SplitMergeFunctionWrapper({ $0 + $1 }))
+        mergeOutputShape: { (l1, l2) in l1 }, mergeFn: { $0 + $1 })
 
-    return convPlusResidual.then(AutoFunction(fnShape: { $0 }, fn: { (prev: Tensor<Float>) in relu(prev) }))
+    return convPlusResidual.then(AutoFunction(fnShape: { $0 }, fn: { relu($0) }))
 }
 
 public typealias AutoResNet = AutoSequenced<AutoSequenced<AutoSequenced<AutoSequenced<AutoSequenced<AutoConvBN, AutoFunction<Tensor<Float>, Tensor<Float>, AutoConv2D<Float>.OutputShape, AutoMaxPool2D<Float>.InputShape>>, AutoMaxPool2D<Float>>, AutoSequencedMany<AutoResidualBlock>>, AutoGlobalAvgPool2D<Float>>, AutoDense<Float>>

Models/LayerInit/AutoSplitMerge.swift

@@ -39,11 +39,15 @@ where Layer1.InputShape == Layer2.InputShape, Layer1.InstanceType.Input == Layer
     public typealias InputShape = Layer1.InputShape
     public typealias OutputShape = OutputShape
 
-    public init(layer1: Layer1, layer2: Layer2, mergeOutputShape: @escaping (Layer1.OutputShape, Layer2.OutputShape) -> OutputShape, mergeFn: SplitMergeFunctionWrapper<Layer1.InstanceType.Output, Layer2.InstanceType.Output, CommonOutput>) {
+    public init(
+        layer1: Layer1, layer2: Layer2,
+        mergeOutputShape: @escaping (Layer1.OutputShape, Layer2.OutputShape) -> OutputShape,
+        mergeFn: @escaping @differentiable (Layer1.InstanceType.Output, Layer2.InstanceType.Output) -> CommonOutput
+    ) {
         self.layer1 = layer1
         self.layer2 = layer2
         self.mergeOutputShape = mergeOutputShape
-        self.mergeFn = mergeFn
+        self.mergeFn = SplitMergeFunctionWrapper(mergeFn)
     }
 
     public func buildModelWithOutputShape<Prefix>(inputShape: Layer1.InputShape, keyPathSoFar: KeyPath<Prefix, InstanceType>, keyDict: inout [AnyAutoLayerKey: Any]) -> (InstanceType, OutputShape) {