struct Dense:  Differentiable {
  var w: Float

  var b: Float

  func forward(_ x: Float) -> Float {
    return w * x + b
  }
}


print("Now gradient for one layer.")
let layer = Dense(w: 1.0, b: 2.0)
let grad = gradient(at: layer) { model in
  model.forward(2.0)
}

print("grad: \(grad)")

struct Model: Differentiable {

  var layer_1 = Dense(w: 1.0, b: 2.0)
  var layer_2 = Dense(w: 3.0, b: 5.0)

  func forward(_ x: Float) -> Float {
    return layer_2.forward(layer_1.forward(x))
  }
}

print("Now gradient for model (nested layers).")
let model = Model()
let gradForModel = gradient(at: model) { model in
  model.forward(2.0)
}

print("grad: \(gradForModel)")