{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import TensorFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public struct MyModel: Layer {\n",
    "    public var layer1: Dense<Float>\n",
    "    public var layer2: Dense<Float>\n",
    "    \n",
    "    public init(nIn: Int, nHid: Int, nOut: Int){\n",
    "        layer1 = Dense(inputSize: nIn, outputSize: nHid)\n",
    "        layer2 = Dense(inputSize: nHid, outputSize: nOut)\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
    "        return input.sequenced(through: layer1, layer2)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var model = MyModel(nIn: 20, nHid: 30, nOut: 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let x: Tensor<Float> = Tensor(randomNormal: TensorShape([10,20]))\n",
    "let y: Tensor<Float> = Tensor(randomNormal: TensorShape([10,10]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let (loss, grads) = model.valueWithGradient { model in \n",
    "    meanAbsoluteError(predicted: model(x), expected: y)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "// FALayer is a layer that supports callbacks through its LayerDelegate.\n",
    "public protocol FALayer: Layer {\n",
    "    var delegates: [(Output) -> ()] { get set }\n",
    "    \n",
    "    // FALayer's will implement this instead of `func callAsFunc`.\n",
    "    @differentiable\n",
    "    func forward(_ input: Input) -> Output\n",
    "    \n",
    "    associatedtype Input\n",
    "    associatedtype Output\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "//export\n",
    "public extension FALayer {\n",
    "    @differentiable(vjp: callGrad)\n",
    "    func callAsFunction(_ input: Input) -> Output {\n",
    "        let activation = forward(input)\n",
    "        for d in delegates { d(activation) }\n",
    "        return activation\n",
    "    }\n",
    "       \n",
    "    // NOTE: AutoDiff synthesizes a leaking VJP for this, so we define a custom VJP.\n",
    "    //    TF-475: https://bugs.swift.org/browse/TF-475\n",
    "    // NOTE: If we use `@differentiating`, then there is a linker error. So we use `@differentiable` instead.\n",
    "    //    TF-476: https://bugs.swift.org/browse/TF-476\n",
    "    func callGrad(_ input: Input) ->\n",
    "        (Output, (Self.Output.TangentVector) -> (Self.TangentVector, Self.Input.TangentVector)) {\n",
    "        return Swift.valueWithPullback(at: self, input) { (m, i) in m(i) }\n",
    "    }\n",
    "    \n",
    "    //We also add a default init to our `delegates` variable, so that we don't have to define it each time, as\n",
    "    //well as a function to easily add a delegate.\n",
    "    var delegates: [(Output) -> ()] { \n",
    "        get { return [] }\n",
    "        set {}\n",
    "    }\n",
    "    \n",
    "    mutating func addDelegate(_ d: @escaping (Output) -> ()) { delegates.append(d) }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "//export\n",
    "\n",
    "@frozen\n",
    "public struct FADense<Scalar: TensorFlowFloatingPoint>: FALayer {\n",
    "    // Note: remove the explicit typealiases after TF-603 is resolved.\n",
    "    public typealias Input = Tensor<Scalar>\n",
    "    public typealias Output = Tensor<Scalar>\n",
    "    public var weight: Tensor<Scalar>\n",
    "    public var bias: Tensor<Scalar>\n",
    "    public typealias Activation = @differentiable (Tensor<Scalar>) -> Tensor<Scalar>\n",
    "    @noDerivative public let activation: Activation\n",
    "\n",
    "    public init(\n",
    "        weight: Tensor<Scalar>,\n",
    "        bias: Tensor<Scalar>,\n",
    "        activation: @escaping Activation\n",
    "    ) {\n",
    "        self.weight = weight\n",
    "        self.bias = bias\n",
    "        self.activation = activation\n",
    "    }\n",
    "\n",
    "    @differentiable\n",
    "    public func forward(_ input: Tensor<Scalar>) -> Tensor<Scalar> {\n",
    "        return activation(input • weight + bias)\n",
    "    }\n",
    "}\n",
    "\n",
    "public extension FADense {\n",
    "    init(_ nIn: Int, _ nOut: Int, activation: @escaping Activation = identity) {\n",
    "        self.init(weight: Tensor(randomNormal: [nIn, nOut]),\n",
    "                  bias: Tensor(zeros: [nOut]),\n",
    "                  activation: activation)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public struct MyModel: Layer {\n",
    "    public var layer1: FADense<Float>\n",
    "    public var layer2: FADense<Float>\n",
    "    \n",
    "    public init(nIn: Int, nHid: Int, nOut: Int){\n",
    "        layer1 = FADense(nIn, nHid, activation: relu)\n",
    "        layer2 = FADense(nHid, nOut)\n",
    "    }\n",
    "    \n",
    "    @differentiable\n",
    "    public func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {\n",
    "        return input.sequenced(through: layer1, layer2)\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "var model = MyModel(nIn: 20, nHid: 30, nOut: 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "let (loss, grads) = model.valueWithGradient { model in \n",
    "    meanAbsoluteError(predicted: model(x), expected: y)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Swift",
   "language": "swift",
   "name": "swift"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}