Structural generic layers 2

[saeta]

Use structural generic programming to build an HParam object.

This change explores using structural generic programming to build (at compile time)
a type-safe "hyper-parameter" object that can subsequently be manipulated and
subsequently used to initialize a NN with the minimum of fuss.

See HParamInitExample.swift for an example usage (partially replicated, and simultaneously idealized here -- there are a few bugs / issues that must be resolved first before we can achieve this):

public struct MyModel: Structural, HParamInitLayer, SequentialLayer {
    var conv: Conv2D<Float>
    var flatten: Flatten<Float>
    var dense: Dense<Float>
}

// Usage:
func makeModel() -> MyModel {
    var hparams = MyModel.HParam()
    hparams.conv = .init(height: 3, width: 3, channels: 10)  // Fully typesafe!
    hparams.dense = .init(size: 10)

    return hparams.build(for: Tensor(zeros: [5, 28, 28, 1]))
}

Note: in order to build the full API, I needed to make some modifications to the proposed
Structural APIs. (Note: this is a quick-hack, and deserves much more refinement!)

For comparisons, writing this out using the existing APIs would look something like the following and require (1) explicit shape calculations, (2) custom definition of hyperparameters, (3) redundant specification of the forward pass + initialization logic:

public struct MyModel: Layer {
  var conv: Conv2D<Float>
  var flatten: Flatten<Float>
  var dense: Dense<Float>

  // TODO: Pass through the other hyperparameters to `conv` and `dense`!
  /// - Precondition: The size of the input to the denseShape must exactly correspond to the size of the input fed to it after running through the convolution. In this case, it will be `floor((inputTensor.shape[1] - convShape.1 + 1) * (inputTensor.shape[2] - convShape.2 + 1) * convShape.4)` as no padding is being used, and strides are (1, 1). (Note: Checked only when trying to actually use the model to perform inference or training).
  /// - Precondition: convShape.3 must be exactly equal to `input.shape[-1]` from the `callAsFunction` parameter.
  public init(convShape: (Int, Int, Int, Int), denseShape: (Int, Int)) {
    conv = Conv2D(convShape)
    flatten = Flatten()
    dense = Dense(inputSize: denseShape.0, outputSize: denseShape.1)
  }

  @differentiable
  func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
    return input.sequenced(through: conv, flatten, dense)
  }
}

Note: for an accurate comparison the initializer should provide a variety of additional hyperparameters, such as strides / etc. The issue with struct & init composition is that default arguments don't compose, and instead must be specified at every level of composition. By representing the information within an explicit alternative type, we sidestep this issue.

Further ideas to explore include:

1. using extensions to add convenience initializers for the HParams objects, such as:

extension MyModel.HParam {
  public enum Flavor {
    case mnist
    case imageNet
  }
  public init(_ flavor: Flavor) {
    self.init()
    self.conv = .init(height: 3, width: 3, channels: 10)
    switch flavor {
    case .mnist:
      self.dense = .init(size: 10)
    case .imageNet:
      self.dense = .init(size: 1000)
    }    
  }
}

with example usage:

var model = MyModel.HParam(.mnist).build(for: Tensor(zeros: [10, 28, 28, 1]))

2. Using a combination of data structures & types to represent skip connections / repeated layers / etc.

Note: of course, we can always fall back to explicitly writing init and func callAsFunction as desired for arbitrarily complicated control flow.

3. Using structural generic programming on the HParam types themselves (e.g. for easy serialization / deserialization, checkpointing, flag parsing, etc).

4. Take an explicit random number generator when building the model from the HParam type to allow random number generation to be made deterministic.

Related PR: #613

[saeta]

CC @shadaj @shabalind @dabrahams

[dabrahams]

Closing until structural GP is available in the upstream compiler.