Merge #1238

1238: Fix inline code block r=dhairyagandhi96 a=harryscholes

### PR Checklist

- [ ] Tests are added
- [ ] Entry in NEWS.md
- [x] Documentation, if applicable
- [ ] Final review from `@MikeInnes` or `@dhairyagandhi96` (for API changes).


Co-authored-by: harryscholes <harryscholes@gmail.com>

.gitattributes

@@ -0,0 +1,2 @@
+paper/* linguist-documentation
+CITATION.bib linguist-detectable=false

.github/FUNDING.yml

@@ -0,0 +1 @@
+custom: https://numfocus.salsalabs.org/donate-to-julia/index.html

.github/pull_request_template.md

@@ -0,0 +1,12 @@
+[Please delete this text and describe your change here.
+For bugfixes, please detail the bug and include a test case which your patch fixes.
+If you are adding a new feature, please clearly describe the design, its rationale, the possible alternatives considered.
+It is easiest to merge new features when there is clear precedent in other systems; we need to know we're taking
+the right direction since it can be hard to change later.]
+
+### PR Checklist
+
+- [ ] Tests are added
+- [ ] Entry in NEWS.md
+- [ ] Documentation, if applicable
+- [ ] Final review from `@MikeInnes` or `@dhairyagandhi96` (for API changes).

.github/workflows/CompatHelper.yml

@@ -0,0 +1,16 @@
+name: CompatHelper
+
+on:
+  schedule:
+    - cron: '00 00 * * *'
+
+jobs:
+  CompatHelper:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Pkg.add("CompatHelper")
+        run: julia -e 'using Pkg; Pkg.add("CompatHelper")'
+      - name: CompatHelper.main()
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: julia -e 'using CompatHelper; CompatHelper.main()'

.github/workflows/TagBot.yml

@@ -0,0 +1,11 @@
+name: TagBot
+on:
+  schedule:
+    - cron: 0 * * * *
+jobs:
+  TagBot:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: JuliaRegistries/TagBot@v1
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }}

.gitignore

@@ -3,5 +3,4 @@
 *.jl.mem
 docs/build/
 docs/site/
-docs/flux.css
-demos
+deps

.gitlab-ci.yml

@@ -0,0 +1,41 @@
+include:
+  - 'https://raw.githubusercontent.com/JuliaGPU/gitlab-ci/master/templates/v6.yml'
+
+image: nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04
+
+
+# julia:1.0:
+#   extends:
+#     - .julia:1.0
+#     - .test
+#   tags:
+#     - nvidia
+# 
+# julia:1.1:
+#   extends:
+#     - .julia:1.1
+#     - .test
+#   tags:
+#     - nvidia
+# 
+# julia:1.2:
+#   extends:
+#     - .julia:1.2
+#     - .test
+#   tags:
+#     - nvidia
+
+julia:1.3:
+  extends:
+    - .julia:1.3
+    - .test
+  tags:
+    - nvidia
+
+julia:nightly:
+  extends:
+    - .julia:nightly
+    - .test
+  tags:
+    - nvidia
+  allow_failure: true

.travis.yml

@@ -1,14 +1,32 @@
 # Documentation: http://docs.travis-ci.com/user/languages/julia/
 language: julia
+
 os:
   - linux
   # - osx
+
 julia:
-  - 0.6
-# uncomment the following lines to override the default test script
+  - 1.3
+  - 1
+  - nightly
+
+notifications:
+  email: false
+
+jobs:
+  include:
+    - stage: "Documentation"
+      julia: 1.3
+      os: linux
+      script:
+        - julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd()));
+                                               Pkg.instantiate()'
+        - julia --project=docs/ docs/make.jl
+      after_success: skip
+
+  allow_failures:
+    - julia: nightly
+
+## uncomment the following lines to override the default test script
 script:
-  - if [[ -a .git/shallow ]]; then git fetch --unshallow; fi
-  - julia -e 'Pkg.clone(pwd()); Pkg.build("Flux"); Pkg.test("Flux"; coverage=true)'
-after_success:
-  - julia -e 'Pkg.add("Documenter")'
-  - julia -e 'cd(Pkg.dir("Flux")); include(joinpath("docs", "make.jl"))'
+ - julia --color=yes -e 'using Pkg; Pkg.activate(); Pkg.instantiate(); Pkg.test()'

CITATION.bib

@@ -0,0 +1,29 @@
+@article{Flux.jl-2018,
+  author    = {Michael Innes and
+               Elliot Saba and
+               Keno Fischer and
+               Dhairya Gandhi and
+               Marco Concetto Rudilosso and
+               Neethu Mariya Joy and
+               Tejan Karmali and
+               Avik Pal and
+               Viral Shah},
+  title     = {Fashionable Modelling with Flux},
+  journal   = {CoRR},
+  volume    = {abs/1811.01457},
+  year      = {2018},
+  url       = {http://arxiv.org/abs/1811.01457},
+  archivePrefix = {arXiv},
+  eprint    = {1811.01457},
+  timestamp = {Thu, 22 Nov 2018 17:58:30 +0100},
+  biburl    = {https://dblp.org/rec/bib/journals/corr/abs-1811-01457},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+
+@article{innes:2018,
+  author    = {Mike Innes},
+  title     = {Flux: Elegant Machine Learning with Julia},
+  journal   = {Journal of Open Source Software},
+  year      = {2018},
+  doi       = {10.21105/joss.00602},
+}

LICENSE.md

@@ -1,6 +1,6 @@
 The Flux.jl package is licensed under the MIT "Expat" License:
 
-> Copyright (c) 2016: Mike Innes.
+> Copyright (c) 2016-19: Julia Computing, INc., Mike Innes and Contributors
 >
 > Permission is hereby granted, free of charge, to any person obtaining
 > a copy of this software and associated documentation files (the

Manifest.toml

@@ -0,0 +1,387 @@
+# This file is machine-generated - editing it directly is not advised
+
+[[AbstractFFTs]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "051c95d6836228d120f5f4b984dd5aba1624f716"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "0.5.0"
+
+[[AbstractTrees]]
+deps = ["Markdown"]
+git-tree-sha1 = "33e450545eaf7699da1a6e755f9ea65f14077a45"
+uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+version = "0.3.3"
+
+[[Adapt]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "fd04049c7dd78cfef0b06cdc1f0f181467655712"
+uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+version = "1.1.0"
+
+[[ArrayLayouts]]
+deps = ["FillArrays", "LinearAlgebra"]
+git-tree-sha1 = "a504dca2ac7eda8761c8f7c1ed52427a1be75a3c"
+uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
+version = "0.2.6"
+
+[[Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[BinaryProvider]]
+deps = ["Libdl", "Logging", "SHA"]
+git-tree-sha1 = "ecdec412a9abc8db54c0efc5548c64dfce072058"
+uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
+version = "0.5.10"
+
+[[CEnum]]
+git-tree-sha1 = "1b77a77c3b28e0b3f413f7567c9bb8dd9bdccd14"
+uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+version = "0.3.0"
+
+[[CUDAapi]]
+deps = ["Libdl", "Logging"]
+git-tree-sha1 = "831b825d10104bd29e28f6da93312a976830717b"
+uuid = "3895d2a7-ec45-59b8-82bb-cfc6a382f9b3"
+version = "4.0.0"
+
+[[CUDAdrv]]
+deps = ["CEnum", "CUDAapi", "Printf"]
+git-tree-sha1 = "f56bbf18c86bcff7a961a32a4947a5abb2963a29"
+uuid = "c5f51814-7f29-56b8-a69c-e4d8f6be1fde"
+version = "6.3.0"
+
+[[CUDAnative]]
+deps = ["Adapt", "BinaryProvider", "CEnum", "CUDAapi", "CUDAdrv", "ExprTools", "GPUCompiler", "LLVM", "Libdl", "Pkg", "Printf"]
+git-tree-sha1 = "ac86db2b05fdfec96b011e25a504ffe7476e8a68"
+uuid = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
+version = "3.1.0"
+
+[[CodeTracking]]
+deps = ["InteractiveUtils", "UUIDs"]
+git-tree-sha1 = "cab4da992adc0a64f63fa30d2db2fd8bec40cab4"
+uuid = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
+version = "0.5.11"
+
+[[CodecZlib]]
+deps = ["TranscodingStreams", "Zlib_jll"]
+git-tree-sha1 = "ded953804d019afa9a3f98981d99b33e3db7b6da"
+uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
+version = "0.7.0"
+
+[[ColorTypes]]
+deps = ["FixedPointNumbers", "Random"]
+git-tree-sha1 = "c73d9cfc2a9d8433dc77f5bff4bddf46b1d78c20"
+uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
+version = "0.10.3"
+
+[[Colors]]
+deps = ["ColorTypes", "FixedPointNumbers", "InteractiveUtils", "Reexport"]
+git-tree-sha1 = "1e9bba7984e78aa8cdeea7f9f7cc984ad4e4b1c7"
+uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
+version = "0.12.2"
+
+[[CommonSubexpressions]]
+deps = ["Test"]
+git-tree-sha1 = "efdaf19ab11c7889334ca247ff4c9f7c322817b0"
+uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
+version = "0.2.0"
+
+[[CompilerSupportLibraries_jll]]
+deps = ["Libdl", "Pkg"]
+git-tree-sha1 = "7c4f882c41faa72118841185afc58a2eb00ef612"
+uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
+version = "0.3.3+0"
+
+[[Cthulhu]]
+deps = ["CodeTracking", "InteractiveUtils", "REPL", "UUIDs", "Unicode"]
+git-tree-sha1 = "f3643e78353199d3097821e806348bd83f364155"
+uuid = "f68482b8-f384-11e8-15f7-abe071a5a75f"
+version = "1.1.1"
+
+[[CuArrays]]
+deps = ["AbstractFFTs", "Adapt", "CEnum", "CUDAapi", "CUDAdrv", "CUDAnative", "DataStructures", "GPUArrays", "Libdl", "LinearAlgebra", "MacroTools", "NNlib", "Pkg", "Printf", "Random", "Reexport", "Requires", "SparseArrays", "Statistics", "TimerOutputs"]
+git-tree-sha1 = "1582b74d2322df7dd94549d4ac9d095e0f20e884"
+uuid = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
+version = "2.2.1"
+
+[[DataAPI]]
+git-tree-sha1 = "176e23402d80e7743fc26c19c681bfb11246af32"
+uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
+version = "1.3.0"
+
+[[DataStructures]]
+deps = ["InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "af6d9c86e191c917c2276fbede1137e8ea20157f"
+uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+version = "0.17.17"
+
+[[Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[DelimitedFiles]]
+deps = ["Mmap"]
+uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+
+[[DiffResults]]
+deps = ["StaticArrays"]
+git-tree-sha1 = "da24935df8e0c6cf28de340b958f6aac88eaa0cc"
+uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+version = "1.0.2"
+
+[[DiffRules]]
+deps = ["NaNMath", "Random", "SpecialFunctions"]
+git-tree-sha1 = "eb0c34204c8410888844ada5359ac8b96292cfd1"
+uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
+version = "1.0.1"
+
+[[Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[ExprTools]]
+git-tree-sha1 = "6f0517056812fd6aa3af23d4b70d5325a2ae4e95"
+uuid = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
+version = "0.1.1"
+
+[[FillArrays]]
+deps = ["LinearAlgebra", "Random", "SparseArrays"]
+git-tree-sha1 = "44f561e293987ffc84272cd3d2b14b0b93123d63"
+uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
+version = "0.8.10"
+
+[[FixedPointNumbers]]
+git-tree-sha1 = "3ba9ea634d4c8b289d590403b4a06f8e227a6238"
+uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
+version = "0.8.0"
+
+[[ForwardDiff]]
+deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "NaNMath", "Random", "SpecialFunctions", "StaticArrays"]
+git-tree-sha1 = "869540e4367122fbffaace383a5bdc34d6e5e5ac"
+uuid = "f6369f11-7733-5829-9624-2563aa707210"
+version = "0.10.10"
+
+[[Functors]]
+deps = ["MacroTools"]
+git-tree-sha1 = "f40adc6422f548176bb4351ebd29e4abf773040a"
+uuid = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+version = "0.1.0"
+
+[[Future]]
+deps = ["Random"]
+uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
+
+[[GPUArrays]]
+deps = ["AbstractFFTs", "Adapt", "LinearAlgebra", "Printf", "Random", "Serialization"]
+git-tree-sha1 = "d887693eb1bd5e1fd573262a978745481895ec7d"
+uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+version = "3.4.1"
+
+[[GPUCompiler]]
+deps = ["Cthulhu", "DataStructures", "InteractiveUtils", "LLVM", "Libdl", "TimerOutputs"]
+git-tree-sha1 = "5275aa268ecd09640b32560e1eae90c78816e4d1"
+uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
+version = "0.2.0"
+
+[[IRTools]]
+deps = ["InteractiveUtils", "MacroTools", "Test"]
+git-tree-sha1 = "90ee39f9beaaa186e4968417ea2b8ed5673c91c0"
+uuid = "7869d1d1-7146-5819-86e3-90919afe41df"
+version = "0.3.3"
+
+[[InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[Juno]]
+deps = ["Base64", "Logging", "Media", "Profile"]
+git-tree-sha1 = "a686b0cf235fa3e491b79b4783c2d2382292b436"
+uuid = "e5e0dc1b-0480-54bc-9374-aad01c23163d"
+version = "0.8.2"
+
+[[LLVM]]
+deps = ["CEnum", "Libdl", "Printf", "Unicode"]
+git-tree-sha1 = "dd3f584c3dbefe39b2a8fbafa1a3b77e31e21255"
+uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
+version = "1.5.1"
+
+[[LibGit2]]
+deps = ["Printf"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[LinearAlgebra]]
+deps = ["Libdl"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[[Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "f7d2e3f654af75f01ec49be82c231c382214223a"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.5"
+
+[[Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[Media]]
+deps = ["MacroTools", "Test"]
+git-tree-sha1 = "75a54abd10709c01f1b86b84ec225d26e840ed58"
+uuid = "e89f7d12-3494-54d1-8411-f7d8b9ae1f27"
+version = "0.5.0"
+
+[[Missings]]
+deps = ["DataAPI"]
+git-tree-sha1 = "de0a5ce9e5289f27df672ffabef4d1e5861247d5"
+uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
+version = "0.4.3"
+
+[[Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[NNlib]]
+deps = ["BinaryProvider", "Libdl", "LinearAlgebra", "Requires", "Statistics"]
+git-tree-sha1 = "d9f196d911f55aeaff11b11f681b135980783824"
+uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+version = "0.6.6"
+
+[[NaNMath]]
+git-tree-sha1 = "928b8ca9b2791081dc71a51c55347c27c618760f"
+uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+version = "0.3.3"
+
+[[OpenSpecFun_jll]]
+deps = ["CompilerSupportLibraries_jll", "Libdl", "Pkg"]
+git-tree-sha1 = "d51c416559217d974a1113522d5919235ae67a87"
+uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
+version = "0.5.3+3"
+
+[[OrderedCollections]]
+git-tree-sha1 = "12ce190210d278e12644bcadf5b21cbdcf225cd3"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.2.0"
+
+[[Pkg]]
+deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+
+[[Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[Profile]]
+deps = ["Printf"]
+uuid = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+
+[[REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[Random]]
+deps = ["Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[Reexport]]
+deps = ["Pkg"]
+git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "0.2.0"
+
+[[Requires]]
+deps = ["UUIDs"]
+git-tree-sha1 = "d37400976e98018ee840e0ca4f9d20baa231dc6b"
+uuid = "ae029012-a4dd-5104-9daa-d747884805df"
+version = "1.0.1"
+
+[[SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+
+[[Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[SortingAlgorithms]]
+deps = ["DataStructures", "Random", "Test"]
+git-tree-sha1 = "03f5898c9959f8115e30bc7226ada7d0df554ddd"
+uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
+version = "0.3.1"
+
+[[SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[SpecialFunctions]]
+deps = ["OpenSpecFun_jll"]
+git-tree-sha1 = "d8d8b8a9f4119829410ecd706da4cc8594a1e020"
+uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
+version = "0.10.3"
+
+[[StaticArrays]]
+deps = ["LinearAlgebra", "Random", "Statistics"]
+git-tree-sha1 = "5c06c0aeb81bef54aed4b3f446847905eb6cbda0"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "0.12.3"
+
+[[Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[StatsBase]]
+deps = ["DataAPI", "DataStructures", "LinearAlgebra", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics"]
+git-tree-sha1 = "a6102b1f364befdb05746f386b67c6b7e3262c45"
+uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+version = "0.33.0"
+
+[[Test]]
+deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
+uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[[TimerOutputs]]
+deps = ["Printf"]
+git-tree-sha1 = "f458ca23ff80e46a630922c555d838303e4b9603"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.6"
+
+[[TranscodingStreams]]
+deps = ["Random", "Test"]
+git-tree-sha1 = "7c53c35547de1c5b9d46a4797cf6d8253807108c"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.9.5"
+
+[[UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[ZipFile]]
+deps = ["Libdl", "Printf", "Zlib_jll"]
+git-tree-sha1 = "254975fef2fc526583bb9b7c9420fe66ffe09f2f"
+uuid = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
+version = "0.9.2"
+
+[[Zlib_jll]]
+deps = ["Libdl", "Pkg"]
+git-tree-sha1 = "a2e0d558f6031002e380a90613b199e37a8565bf"
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+version = "1.2.11+10"
+
+[[Zygote]]
+deps = ["AbstractFFTs", "ArrayLayouts", "DiffRules", "FillArrays", "ForwardDiff", "Future", "IRTools", "InteractiveUtils", "LinearAlgebra", "MacroTools", "NNlib", "NaNMath", "Random", "Requires", "SpecialFunctions", "Statistics", "ZygoteRules"]
+git-tree-sha1 = "707ceea58e2bd0ff3077ab13a92f8355181d3ee4"
+uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
+version = "0.4.20"
+
+[[ZygoteRules]]
+deps = ["MacroTools"]
+git-tree-sha1 = "b3b4882cc9accf6731a08cc39543fbc6b669dca8"
+uuid = "700de1a5-db45-46bc-99cf-38207098b444"
+version = "0.2.0"

NEWS.md

@@ -0,0 +1,61 @@
+# v0.11
+* Change to `DataLoader`'s constructor [https://github.com/FluxML/Flux.jl/pull/1152]
+* Use `DataLoader` with `NamedTuple`s, so that tensors can be accessed by name [https://github.com/FluxML/Flux.jl/pull/1221].
+* Error if Dense layers weights and biases are not arrays [https://github.com/FluxML/Flux.jl/pull/1218].
+
+# v0.10.5
+* Add option for [same padding](https://github.com/FluxML/Flux.jl/pull/901) to conv and pooling layers by setting `pad=SamePad()`.
+* Added option to set `bias` to [Flux.Zeros](https://github.com/FluxML/Flux.jl/pull/873) to eliminating `bias` from being trained.
+* Added `GlobalMaxPool` and `GlobalMeanPool` [layers](https://github.com/FluxML/Flux.jl/pull/950) for performing global pooling operations.
+* Added `ClipValue` and `ClipNorm` in this [pr](https://github.com/FluxML/Flux.jl/pull/1133) to `Flux.Optimise` to provide a cleaner API for gradient clipping.
+* Added new kwarg-only [constructors](https://github.com/FluxML/Flux.jl/pull/873) for the various convolutional layers.
+* Documented the convolutional layer constructors accepting `weight` and `bias` keyword arguments to supply custom arrays for those fields.
+* Testing suite improvements now test for gradients of all layers along with GPU support.
+* Functors have now moved to [Functors.jl](https://github.com/FluxML/Flux.jl/pull/1174) to allow for their use outside of Flux.
+* Added [helper functions](https://github.com/FluxML/Flux.jl/pull/873) `Flux.convfilter` and `Flux.depthwiseconvfilter` to construct weight arrays for convolutions outside of layer constructors so as to not have to depend on the default layers for custom implementations.
+
+# v0.10.0
+* The default AD engine has switched from [Tracker to Zygote.jl](https://github.com/FluxML/Flux.jl/pull/669)
+  - The dependency on Tracker.jl has been removed.
+  - This means Flux now does not depend on using a specialised `TrackedArray` type, and can be used with normal Array implementations directly.
+  - Tracker compatibility is maintained in most common cases, but Zygote will be the preferred AD backend for Flux from now on.
+* The CUDNN wrappers have been [moved from Flux into CuArrays](https://github.com/FluxML/Flux.jl/pull/874), to allow for better supporting the CUDA backend, and improve user experience, not to mention making Flux lean.
+* `*crossentropy` functions now [work as expected with CuArrays](https://github.com/FluxML/Flux.jl/pull/926). [PR for binarycrossentropy](https://github.com/FluxML/Flux.jl/pull/940).
+* Added [clearer docs](https://github.com/FluxML/Flux.jl/pull/904) around training and the Optimiser interface.
+* [Layer initialisations](https://github.com/FluxML/Flux.jl/pull/937) have been improved with a clearer API on how to extend it for other purposes.
+* [Better messaging around CUDA availability](https://github.com/FluxML/Flux.jl/pull/924), with hooks to initialize the GPU as default where possible.
+* `@treelike` has been formalised as a [functor](https://github.com/FluxML/Flux.jl/pull/865), with an effective deprecation.
+* `testmode!` is deprecated in favour of [istraining](https://github.com/FluxML/Flux.jl/pull/669)
+
+# v0.9.0
+* [Depthwise convolutional layer API changes](https://github.com/FluxML/Flux.jl/pull/756) from `in => mult` channel specification to `in => out` channel specification, and deprecates implicit `out` constructor.
+* New [SkipConnection](https://github.com/FluxML/Flux.jl/pull/446), which can be used to train residual neural network architectures.
+* New [RADAM](https://github.com/FluxML/Flux.jl/pull/842) optimiser.
+
+# v0.8.0
+
+* [Dropout now has a `dims` argument for specifying the unbroadcast dimensions.](https://github.com/FluxML/Flux.jl/pull/563)
+* New [ConvTranspose layer](https://github.com/FluxML/Flux.jl/pull/311).
+* New [Maxout layer](https://github.com/FluxML/Flux.jl/pull/647)
+* Datasets are now [hash verified on download](https://github.com/FluxML/Flux.jl/pull/585) to avoid corruption.
+* We now [zero the initial state for RNNs](https://github.com/FluxML/Flux.jl/pull/590/).
+* [Normalisation can now work on arbitrary `dims`.](https://github.com/FluxML/Flux.jl/pull/592)
+* Many docs and bugfixes thanks to @KristofferC and others.
+* [NamedTuples now work like Tuples](https://github.com/FluxML/Flux.jl/pull/603) when doing `mapleaves`.
+* New "performance tips" [section of the docs](https://github.com/FluxML/Flux.jl/pull/615).
+* The training loop is [now more readable](https://github.com/FluxML/Flux.jl/pull/651) and better shows how to use the lower-level APIs.
+* New [AlphaDropout](https://github.com/FluxML/Flux.jl/pull/656).
+* [Data.Iris](https://github.com/FluxML/Flux.jl/pull/652) makes Fisher's Iris dataset available with `Iris.labels` and `Iris.features`.
+* New [InstanceNorm](https://github.com/FluxML/Flux.jl/pull/634), as popularized by [Instance Normalization: The Missing Ingredient for Fast Stylization](https://arxiv.org/abs/1607.08022).
+* New [GroupNorm](https://github.com/FluxML/Flux.jl/pull/696), as described in [Group Normalization](https://arxiv.org/abs/1803.08494).
+* New [CrossCor](https://github.com/FluxML/Flux.jl/pull/762).
+
+AD Changes:
+
+* `det`, `logdet` and `logabsdet` [now have adjoints](https://github.com/FluxML/Flux.jl/pull/596/files).
+* Support for [PermuteDimsArray](https://github.com/FluxML/Flux.jl/pull/576).
+* Flux.Tracker is now its [own package](https://github.com/FluxML/Tracker.jl), in preparation for replacing it with Zygote.
+
+# v0.7.0
+
+Despite the heroic efforts of scholars and archeologists, pre-0.7 history is lost to the sands of time.

Project.toml

@@ -0,0 +1,51 @@
+name = "Flux"
+uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+version = "0.11.0-DEV"
+
+[deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
+Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
+CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
+DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Juno = "e5e0dc1b-0480-54bc-9374-aad01c23163d"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+ZipFile = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+AbstractTrees = "0.2, 0.3"
+Adapt = "1, 2.0"
+CodecZlib = "0.5, 0.6, 0.7"
+Colors = "0.8, 0.9, 0.10, 0.11, 0.12"
+CuArrays = "2"
+Functors = "0.1"
+Juno = "0.5, 0.6, 0.7, 0.8"
+MacroTools = "0.3, 0.4, 0.5"
+NNlib = "0.6"
+Reexport = "0.2"
+StatsBase = "0"
+ZipFile = "0.7, 0.8, 0.9"
+Zygote = "0.4.13"
+julia = "1.3"
+
+[extras]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "Documenter", "IterTools", "LinearAlgebra"]

README.md

@@ -1,11 +1,15 @@
-# Флукс
+<p align="center">
+<img width="400px" src="https://raw.githubusercontent.com/FluxML/fluxml.github.io/master/logo.png"/>
+</p>
 
-[![Build Status](https://travis-ci.org/FluxML/Flux.jl.svg?branch=master)](https://travis-ci.org/FluxML/Flux.jl) [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://fluxml.github.io/Flux.jl/stable/) [![Join the chat at https://gitter.im/FluxML](https://badges.gitter.im/FluxML/Lobby.svg)](https://gitter.im/FluxML/Lobby) [Slack](https://discourse.julialang.org/t/announcing-a-julia-slack/4866)
+[![Build Status](https://travis-ci.org/FluxML/Flux.jl.svg?branch=master)](https://travis-ci.org/FluxML/Flux.jl) [![](https://img.shields.io/badge/docs-stable-blue.svg)](https://fluxml.github.io/Flux.jl/stable/) [![](https://img.shields.io/badge/chat-on%20slack-yellow.svg)](https://slackinvite.julialang.org/) [![DOI](https://joss.theoj.org/papers/10.21105/joss.00602/status.svg)](https://doi.org/10.21105/joss.00602)
 
-Flux is a refreshing approach to machine learning. It provides lightweight abstractions on top of Julia's native GPU and AD support, while remaining fully hackable (right down to the [GPU kernels](https://github.com/FluxML/CuArrays.jl)).
+Flux is an elegant approach to machine learning. It's a 100% pure-Julia stack, and provides lightweight abstractions on top of Julia's native GPU and AD support. Flux makes the easy things easy while remaining fully hackable.
 
 ```julia
-julia> Pkg.add("Flux")
+] add Flux
 ```
 
-See the [documentation](http://fluxml.github.io/Flux.jl/stable/) or the [model zoo](https://github.com/FluxML/model-zoo/) for examples.
+See the [documentation](https://fluxml.github.io/Flux.jl/) or the [model zoo](https://github.com/FluxML/model-zoo/) for examples.
+
+If you use Flux in your research, please [cite](CITATION.bib) our work.

REQUIRE

@@ -1,7 +0,0 @@
-julia 0.6.0
-DataFlow 0.2.1
-Juno
-MacroTools 0.3.3
-NNlib
-ForwardDiff
-Requires

bors.toml

@@ -0,0 +1,4 @@
+status = [
+  "ci/gitlab%"
+]
+timeout-sec = 7200

docs/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+
+[compat]
+Documenter = "0.24"

docs/make.jl

@@ -1,27 +1,36 @@
-using Documenter, Flux
+using Documenter, Flux, NNlib
 
-makedocs(modules=[Flux],
-         doctest = false,
-         format = :html,
-         analytics = "UA-36890222-9",
+DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive=true)
+makedocs(modules=[Flux, NNlib],
+         doctest = VERSION >= v"1.4",
          sitename = "Flux",
-         assets = ["../flux.css"],
          pages = ["Home" => "index.md",
                   "Building Models" =>
                     ["Basics" => "models/basics.md",
                      "Recurrence" => "models/recurrence.md",
-                     "Layer Reference" => "models/layers.md"],
+                     "Regularisation" => "models/regularisation.md",
+                     "Model Reference" => "models/layers.md",
+                     "Advanced Model Building" => "models/advanced.md",
+                     "NNlib" => "models/nnlib.md"],
+                  "Handling Data" =>
+                    ["One-Hot Encoding" => "data/onehot.md",
+                     "DataLoader" => "data/dataloader.md"],
                   "Training Models" =>
                     ["Optimisers" => "training/optimisers.md",
                      "Training" => "training/training.md"],
-                  "One-Hot Encoding" => "data/onehot.md",
                   "GPU Support" => "gpu.md",
-                  "Contributing & Help" => "contributing.md"])
+                  "Saving & Loading" => "saving.md",
+                  "The Julia Ecosystem" => "ecosystem.md",
+                  "Utility Functions" => "utilities.md",
+                  "Performance Tips" => "performance.md",
+                  "Datasets" => "datasets.md",
+                  "Community" => "community.md"],
+         format = Documenter.HTML(
+             analytics = "UA-36890222-9",
+             assets = ["assets/flux.css"],
+             prettyurls = get(ENV, "CI", nothing) == "true"),
+         )
 
-deploydocs(
-   repo = "github.com/FluxML/Flux.jl.git",
-   target = "build",
-   osname = "linux",
-   julia = "0.6",
-   deps = nothing,
-   make = nothing)
+deploydocs(repo = "github.com/FluxML/Flux.jl.git",
+           target = "build",
+           push_preview = true)

docs/src/assets/flux.css

@@ -0,0 +1,113 @@
+@import url('https://fonts.googleapis.com/css?family=Lato:400,400i');
+
+body {
+  font-family: Lato, "Segoe UI",Roboto,"Helvetica Neue",Arial,sans-serif;
+}
+
+nav.toc {
+  padding-top: 0;
+  background: rgb(240, 240, 240);
+  line-height: 2em;
+  cursor: default;
+  user-select: none;
+}
+
+h1+h2 {
+  margin-top: 0;
+}
+
+/* Green banner in ToC */
+nav.toc > h1 {
+  margin-top: 0;
+  padding-top: 0.4em;
+  padding-bottom: 0.5em;
+  border-bottom: 5px solid white;
+  box-shadow: 0px -2px 5px rgb(60,60,60);
+  margin-bottom: 0.5em;
+  background: rgb(60, 150, 60);
+
+  font-style: italic;
+  font-weight: normal;
+  font-size: 50pt;
+  text-transform: lowercase;
+  text-shadow: 2px 2px 5px rgba(0,0,0,0.2);
+  color: white;
+}
+
+/* Reduce ToC font size */
+.toctext {
+  font-size: 10pt;
+}
+
+/* Fade out non-clickable ToC headers */
+nav.toc ul span.toctext {
+  color: rgb(180, 180, 180);
+}
+
+nav.toc ul .toctext {
+  color: rgb(100, 100, 100);
+}
+
+nav.toc ul a.toctext:hover {
+  color: inherit;
+  background: rgb(220, 220, 220);
+  cursor: default;
+}
+
+nav.toc li.current > .toctext {
+  background: linear-gradient(90deg, rgb(245,245,245) 0%, white 90%);
+  font-weight: normal;
+}
+
+nav.toc ul.internal li.toplevel {
+  font-weight: normal;
+}
+
+/* Content */
+
+article { max-width: none; }
+
+article > p, article > ul {
+  max-width: 45em;
+}
+
+/* Links */
+a, a:visited { color: rgb(0, 120, 0); }
+article p a { border-bottom: 1px solid rgb(200, 230, 200); }
+a:hover, a:visited:hover { color: rgb(0, 80, 0); }
+
+/* Article Links */
+article p a { border-bottom: 1px solid rgb(200, 230, 200); }
+article p a:hover, article a:visited:hover { color: rgb(0, 120, 0); }
+article p a:hover { border-bottom: 1px solid rgb(150, 200, 150); }
+
+/* Doctstrings */
+article section.docstring {
+  padding: 0.5em 0;
+  border-left: none;
+  border-right: none;
+  border-bottom: none;
+}
+
+/* Code */
+
+article pre, article p > code {
+  background: rgb(245, 250, 245);
+}
+
+article pre {
+  border: none;
+  max-width: none;
+  padding: 1em;
+  border-radius: 10px 0px 0px 10px;
+  margin-left: -1em;
+  margin-right: -2em;
+}
+
+.hljs-comment {
+  font-style: italic;
+}
+
+.hljs-number {
+  color: rgb(0, 150, 150);
+}

docs/src/community.md

@@ -0,0 +1,5 @@
+# Community
+
+All Flux users are welcome to join our community on the [Julia forum](https://discourse.julialang.org/), or the [slack](https://discourse.julialang.org/t/announcing-a-julia-slack/4866) (channel #machine-learning). If you have questions or issues we'll try to help you out.
+
+If you're interested in hacking on Flux, the [source code](https://github.com/FluxML/Flux.jl) is open and easy to understand -- it's all just the same Julia code you work with normally. You might be interested in our [intro issues](https://github.com/FluxML/Flux.jl/issues?q=is%3Aopen+is%3Aissue+label%3A%22help+wanted%22) to get started.

docs/src/contributing.md

@@ -1,9 +0,0 @@
-# Contributing & Help
-
-If you need help, please ask on the [Julia forum](https://discourse.julialang.org/), the [slack](https://discourse.julialang.org/t/announcing-a-julia-slack/4866) (channel #machine-learning), or Flux's [Gitter](https://gitter.im/FluxML/Lobby).
-
-Right now, the best way to help out is to try out the examples and report any issues or missing features as you find them. The second best way is to help us spread the word, perhaps by [starring the repo](https://github.com/MikeInnes/Flux.jl).
-
-If you're interested in hacking on Flux, most of the [code](https://github.com/MikeInnes/Flux.jl/tree/master/src) is pretty straightforward. Adding new [layer definitions](https://github.com/MikeInnes/Flux.jl/tree/master/src/layers) or cost functions is simple using the Flux DSL itself, and things like data utilities and training processes are all plain Julia code.
-
-If you get stuck or need anything, let us know!

docs/src/data/dataloader.md

@@ -0,0 +1,6 @@
+# DataLoader
+Flux provides the `DataLoader` type in the `Flux.Data` module to handle iteration over mini-batches of data. 
+
+```@docs
+Flux.Data.DataLoader
+```

docs/src/data/onehot.md

@@ -3,37 +3,42 @@
 It's common to encode categorical variables (like `true`, `false` or `cat`, `dog`) in "one-of-k" or ["one-hot"](https://en.wikipedia.org/wiki/One-hot) form. Flux provides the `onehot` function to make this easy.
 
 ```
-julia> using Flux: onehot
+julia> using Flux: onehot, onecold
 
 julia> onehot(:b, [:a, :b, :c])
 3-element Flux.OneHotVector:
- false
-  true
- false
+ 0
+ 1
+ 0
 
 julia> onehot(:c, [:a, :b, :c])
 3-element Flux.OneHotVector:
- false
- false
-  true
+ 0
+ 0
+ 1
 ```
 
-The inverse is `argmax` (which can take a general probability distribution, as well as just booleans).
+The inverse is `onecold` (which can take a general probability distribution, as well as just booleans).
 
 ```julia
-julia> argmax(ans, [:a, :b, :c])
+julia> onecold(ans, [:a, :b, :c])
 :c
 
-julia> argmax([true, false, false], [:a, :b, :c])
+julia> onecold([true, false, false], [:a, :b, :c])
 :a
 
-julia> argmax([0.3, 0.2, 0.5], [:a, :b, :c])
+julia> onecold([0.3, 0.2, 0.5], [:a, :b, :c])
 :c
 ```
 
+```@docs
+Flux.onehot
+Flux.onecold
+```
+
 ## Batches
 
-`onehotbatch` creates a batch (matrix) of one-hot vectors, and `argmax` treats matrices as batches.
+`onehotbatch` creates a batch (matrix) of one-hot vectors, and `onecold` treats matrices as batches.
 
 ```julia
 julia> using Flux: onehotbatch
@@ -52,3 +57,7 @@ julia> onecold(ans, [:a, :b, :c])
 ```
 
 Note that these operations returned `OneHotVector` and `OneHotMatrix` rather than `Array`s. `OneHotVector`s behave like normal vectors but avoid any unnecessary cost compared to using an integer index directly. For example, multiplying a matrix with a one-hot vector simply slices out the relevant row of the matrix under the hood.
+
+```@docs
+Flux.onehotbatch
+```

docs/src/datasets.md

@@ -0,0 +1,20 @@
+# Datasets
+
+Flux includes several standard machine learning datasets.
+
+```@docs
+Flux.Data.Iris.features()
+Flux.Data.Iris.labels()
+Flux.Data.MNIST.images()
+Flux.Data.MNIST.labels()
+Flux.Data.FashionMNIST.images()
+Flux.Data.FashionMNIST.labels()
+Flux.Data.CMUDict.phones()
+Flux.Data.CMUDict.symbols()
+Flux.Data.CMUDict.rawdict()
+Flux.Data.CMUDict.cmudict()
+Flux.Data.Sentiment.train()
+Flux.Data.Sentiment.test()
+Flux.Data.Sentiment.dev()
+```
+

docs/src/ecosystem.md

@@ -0,0 +1,21 @@
+# The Julia Ecosystem
+
+One of the main strengths of Julia lies in an ecosystem of packages 
+globally providing a rich and consistent user experience.
+
+This is a non-exhaustive list of Julia packages, nicely complementing `Flux` in typical
+machine learning and deep learning workflows:
+
+- [ArgParse.jl](https://github.com/carlobaldassi/ArgParse.jl): package for parsing command-line arguments to Julia programs.
+- [Augmentor.jl](https://github.com/Evizero/Augmentor.jl): a fast image augmentation library in Julia for machine learning.
+- [BSON.jl](https://github.com/JuliaIO/BSON.jl): package for working with the Binary JSON serialisation format
+- [DataFrames.jl](https://github.com/joshday/OnlineStats.jl): in-memory tabular data in Julia
+- [DrWatson.jl](https://github.com/JuliaDynamics/DrWatson.jl):  a scientific project assistant software
+- [MLDatasets.jl](https://github.com/JuliaML/MLDatasets.jl): utility package for accessing common machine learning datasets
+- [OnlineStats.jl](https://github.com/joshday/OnlineStats.jl): single-pass algorithms for statistics
+- [Parameters.jl](https://github.com/mauro3/Parameters.jl): types with default field values, keyword constructors and (un-)pack macros
+- [ProgressMeters.jl](https://github.com/timholy/ProgressMeter.jl): progress meters for long-running computations
+- [TensorBoardLogger.jl](https://github.com/PhilipVinc/TensorBoardLogger.jl): easy peasy logging to [tensorboard](https://www.tensorflow.org/tensorboard) in Julia
+
+
+This tight integration among Julia pakages is shown in some of the examples in the [model-zoo](https://github.com/FluxML/model-zoo) repository.

docs/src/gpu.md

@@ -1,9 +1,15 @@
 # GPU Support
 
-Support for array operations on other hardware backends, like GPUs, is provided by external packages like [CuArrays](https://github.com/JuliaGPU/CuArrays.jl) and [CLArrays](https://github.com/JuliaGPU/CLArrays.jl). Flux doesn't care what array type you use, so we can just plug these in without any other changes.
+NVIDIA GPU support should work out of the box on systems with CUDA and CUDNN installed. For more details see the [CuArrays](https://github.com/JuliaGPU/CuArrays.jl) readme.
+
+## GPU Usage
+
+Support for array operations on other hardware backends, like GPUs, is provided by external packages like [CuArrays](https://github.com/JuliaGPU/CuArrays.jl). Flux is agnostic to array types, so we simply need to move model weights and data to the GPU and Flux will handle it.
 
 For example, we can use `CuArrays` (with the `cu` converter) to run our [basic example](models/basics.md) on an NVIDIA GPU.
 
+(Note that you need to have CUDA available to use CuArrays – please see the [CuArrays.jl](https://github.com/JuliaGPU/CuArrays.jl) instructions for more details.)
+
 ```julia
 using CuArrays
 
@@ -19,17 +25,52 @@ loss(x, y) # ~ 3
 
 Note that we convert both the parameters (`W`, `b`) and the data set (`x`, `y`) to cuda arrays. Taking derivatives and training works exactly as before.
 
-If you define a structured model, like a `Dense` layer or `Chain`, you just need to convert the internal parameters. Flux provides `mapleaves`, which allows you to alter all parameters of a model at once.
+If you define a structured model, like a `Dense` layer or `Chain`, you just need to convert the internal parameters. Flux provides `fmap`, which allows you to alter all parameters of a model at once.
 
 ```julia
 d = Dense(10, 5, σ)
-d = mapleaves(cu, d)
-d.W # Tracked CuArray
+d = fmap(cu, d)
+d.W # CuArray
 d(cu(rand(10))) # CuArray output
 
 m = Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
-m = mapleaves(cu, m)
+m = fmap(cu, m)
 d(cu(rand(10)))
 ```
 
-The [mnist example](https://github.com/FluxML/model-zoo/blob/master/mnist/mnist.jl) contains the code needed to run the model on the GPU; just uncomment the lines after `using CuArrays`.
+As a convenience, Flux provides the `gpu` function to convert models and data to the GPU if one is available. By default, it'll do nothing, but loading `CuArrays` will cause it to move data to the GPU instead.
+
+```julia
+julia> using Flux, CuArrays
+
+julia> m = Dense(10,5) |> gpu
+Dense(10, 5)
+
+julia> x = rand(10) |> gpu
+10-element CuArray{Float32,1}:
+ 0.800225
+ ⋮
+ 0.511655
+
+julia> m(x)
+5-element CuArray{Float32,1}:
+ -0.30535
+ ⋮
+ -0.618002
+```
+
+The analogue `cpu` is also available for moving models and data back off of the GPU.
+
+```julia
+julia> x = rand(10) |> gpu
+10-element CuArray{Float32,1}:
+ 0.235164
+ ⋮
+ 0.192538
+
+julia> x |> cpu
+10-element Array{Float32,1}:
+ 0.235164
+ ⋮
+ 0.192538
+```

docs/src/index.md

@@ -1,14 +1,17 @@
 # Flux: The Julia Machine Learning Library
 
-Flux is a library for machine learning. It comes "batteries-included" with many useful tools built in, but also lets you use the full power of the Julia language where you need it. The whole stack is implemented in clean Julia code (right down to the [GPU kernels](https://github.com/FluxML/CuArrays.jl)) and any part can be tweaked to your liking.
+Flux is a library for machine learning. It comes "batteries-included" with many useful tools built in, but also lets you use the full power of the Julia language where you need it. We follow a few key principles:
 
-# Installation
+* **Doing the obvious thing**. Flux has relatively few explicit APIs for features like regularisation or embeddings. Instead, writing down the mathematical form will work – and be fast.
+* **You could have written Flux**. All of it, from [LSTMs](https://github.com/FluxML/Flux.jl/blob/ec16a2c77dbf6ab8b92b0eecd11661be7a62feef/src/layers/recurrent.jl#L131) to [GPU kernels](https://github.com/JuliaGPU/CuArrays.jl), is straightforward Julia code. When in doubt, it’s well worth looking at [the source](https://github.com/FluxML/Flux.jl/). If you need something different, you can easily roll your own.
+* **Play nicely with others**. Flux works well with Julia libraries from [data frames](https://github.com/JuliaComputing/JuliaDB.jl) and [images](https://github.com/JuliaImages/Images.jl) to [differential equation solvers](https://github.com/JuliaDiffEq/DifferentialEquations.jl), so you can easily build complex data processing pipelines that integrate Flux models.
 
-Install [Julia 0.6.0 or later](https://julialang.org/downloads/), if you haven't already.
+## Installation
 
-```julia
-Pkg.add("Flux")
-Pkg.test("Flux") # Check things installed correctly
-```
+Download [Julia 1.0](https://julialang.org/) or later, if you haven't already. You can add Flux from using Julia's package manager, by typing `] add Flux` in the Julia prompt.
 
-Start with the [basics](models/basics.md). The [model zoo](https://github.com/FluxML/model-zoo/) is also a good starting point for many common kinds of models.
+If you have CUDA you can also run `] add CuArrays` to get GPU support; see [here](gpu.md) for more details.
+
+## Learning Flux
+
+There are several different ways to learn Flux. If you just want to get started writing models, the [model zoo](https://github.com/FluxML/model-zoo/) gives good starting points for many common ones. This documentation provides a reference to all of Flux's APIs, as well as a from-scratch introduction to Flux's take on models and how they work. Once you understand these docs, congratulations, you also understand [Flux's source code](https://github.com/FluxML/Flux.jl), which is intended to be concise, legible and a good reference for more advanced concepts.

docs/src/models/advanced.md

@@ -0,0 +1,73 @@
+# Advanced Model Building and Customisation
+
+Here we will try and describe usage of some more advanced features that Flux provides to give more control over model building.
+
+## Customising Parameter Collection for a Model
+
+Taking reference from our example `Affine` layer from the [basics](basics.md#Building-Layers-1).
+
+By default all the fields in the `Affine` type are collected as its parameters, however, in some cases it may be desired to hold other metadata in our "layers" that may not be needed for training, and are hence supposed to be ignored while the parameters are collected. With Flux, it is possible to mark the fields of our layers that are trainable in two ways.
+
+The first way of achieving this is through overloading the `trainable` function.
+
+```julia-repl
+julia> @functor Affine
+
+julia> a = Affine(rand(3,3), rand(3))
+Affine{Array{Float64,2},Array{Float64,1}}([0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297], [0.42394, 0.0170927, 0.544955])
+
+julia> Flux.params(a) # default behavior
+Params([[0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297], [0.42394, 0.0170927, 0.544955]])
+
+julia> Flux.trainable(a::Affine) = (a.W,)
+
+julia> Flux.params(a)
+Params([[0.66722 0.774872 0.249809; 0.843321 0.403843 0.429232; 0.683525 0.662455 0.065297]])
+```
+
+Only the fields returned by `trainable` will be collected as trainable parameters of the layer when calling `Flux.params`.
+
+Another way of achieving this is through the `@functor` macro directly. Here, we can mark the fields we are interested in by grouping them in the second argument:
+
+```julia
+Flux.@functor Affine (W,)
+```
+
+However, doing this requires the `struct` to have a corresponding constructor that accepts those parameters.
+
+## Freezing Layer Parameters
+
+When it is desired to not include all the model parameters (for e.g. transfer learning), we can simply not pass in those layers into our call to `params`.
+
+Consider a simple multi-layer perceptron model where we want to avoid optimising the first two `Dense` layers. We can obtain
+this using the slicing features `Chain` provides:
+
+```julia
+m = Chain(
+      Dense(784, 64, relu),
+      Dense(64, 64, relu),
+      Dense(32, 10)
+    )
+
+ps = Flux.params(m[3:end])
+```
+
+The `Zygote.Params` object `ps` now holds a reference to only the parameters of the layers passed to it.
+
+During training, the gradients will only be computed for (and applied to) the last `Dense` layer, therefore only that would have its parameters changed.
+
+`Flux.params` also takes multiple inputs to make it easy to collect parameters from heterogenous models with a single call. A simple demonstration would be if we wanted to omit optimising the second `Dense` layer in the previous example. It would look something like this:
+
+```julia
+Flux.params(m[1], m[3:end])
+```
+
+Sometimes, a more fine-tuned control is needed. 
+We can freeze a specific parameter of a specific layer which already entered a `Params` object `ps`, 
+by simply deleting it from `ps`:
+
+```julia
+ps = params(m)
+delete!(ps, m[2].b) 
+```
+

docs/src/models/basics.md

@@ -2,58 +2,114 @@
 
 ## Taking Gradients
 
-Consider a simple linear regression, which tries to predict an output array `y` from an input `x`. (It's a good idea to follow this example in the Julia repl.)
+Flux's core feature is taking gradients of Julia code. The `gradient` function takes another Julia function `f` and a set of arguments, and returns the gradient with respect to each argument. (It's a good idea to try pasting these examples in the Julia terminal.)
+
+```jldoctest basics
+julia> using Flux
+
+julia> f(x) = 3x^2 + 2x + 1;
+
+julia> df(x) = gradient(f, x)[1]; # df/dx = 6x + 2
+
+julia> df(2)
+14
+
+julia> d2f(x) = gradient(df, x)[1]; # d²f/dx² = 6
+
+julia> d2f(2)
+6
+```
+
+When a function has many parameters, we can get gradients of each one at the same time:
+
+```jldoctest basics
+julia> f(x, y) = sum((x .- y).^2);
+
+julia> gradient(f, [2, 1], [2, 0])
+([0, 2], [0, -2])
+```
+
+But machine learning models can have *hundreds* of parameters! To handle this, Flux lets you work with collections of parameters, via `params`. You can get the gradient of all parameters used in a program without explicitly passing them in.
+
+```jldoctest basics
+julia> x = [2, 1];
+
+julia> y = [2, 0];
+
+julia> gs = gradient(params(x, y)) do
+         f(x, y)
+       end
+Grads(...)
+
+julia> gs[x]
+2-element Array{Int64,1}:
+ 0
+ 2
+
+julia> gs[y]
+2-element Array{Int64,1}:
+  0
+ -2
+```
+
+Here, `gradient` takes a zero-argument function; no arguments are necessary because the `params` tell it what to differentiate.
+
+This will come in really handy when dealing with big, complicated models. For now, though, let's start with something simple.
+
+## Simple Models
+
+Consider a simple linear regression, which tries to predict an output array `y` from an input `x`.
 
 ```julia
 W = rand(2, 5)
 b = rand(2)
 
 predict(x) = W*x .+ b
-loss(x, y) = sum((predict(x) .- y).^2)
+
+function loss(x, y)
+  ŷ = predict(x)
+  sum((y .- ŷ).^2)
+end
 
 x, y = rand(5), rand(2) # Dummy data
 loss(x, y) # ~ 3
 ```
 
-To improve the prediction we can take the gradients of `W` and `b` with respect to the loss function and perform gradient descent. We could calculate gradients by hand, but Flux will do it for us if we tell it that `W` and `b` are trainable *parameters*.
+To improve the prediction we can take the gradients of `W` and `b` with respect to the loss and perform gradient descent.
 
 ```julia
-using Flux.Tracker: param, back!, data, grad
+using Flux
 
-W = param(W)
-b = param(b)
-
-l = loss(x, y)
-
-back!(l)
+gs = gradient(() -> loss(x, y), params(W, b))
 ```
 
-`loss(x, y)` returns the same number, but it's now a *tracked* value that records gradients as it goes along. Calling `back!` then calculates the gradient of `W` and `b`. We can see what this gradient is, and modify `W` to train the model.
+Now that we have gradients, we can pull them out and update `W` to train the model.
 
 ```julia
-grad(W)
+W̄ = gs[W]
 
-# Update the parameter
-W.data .-= 0.1grad(W)
+W .-= 0.1 .* W̄
 
 loss(x, y) # ~ 2.5
 ```
 
 The loss has decreased a little, meaning that our prediction `x` is closer to the target `y`. If we have some data we can already try [training the model](../training/training.md).
 
-All deep learning in Flux, however complex, is a simple generalisation of this example. Of course, models can *look* very different – they might have millions of parameters or complex control flow, and there are ways to manage this complexity. Let's see what that looks like.
+All deep learning in Flux, however complex, is a simple generalisation of this example. Of course, models can *look* very different – they might have millions of parameters or complex control flow. Let's see how Flux handles more complex models.
 
 ## Building Layers
 
 It's common to create more complex models than the linear regression above. For example, we might want to have two linear layers with a nonlinearity like [sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function) (`σ`) in between them. In the above style we could write this as:
 
 ```julia
-W1 = param(rand(3, 5))
-b1 = param(rand(3))
+using Flux
+
+W1 = rand(3, 5)
+b1 = rand(3)
 layer1(x) = W1 * x .+ b1
 
-W2 = param(rand(2, 3))
-b2 = param(rand(2))
+W2 = rand(2, 3)
+b2 = rand(2)
 layer2(x) = W2 * x .+ b2
 
 model(x) = layer2(σ.(layer1(x)))
@@ -65,8 +121,8 @@ This works but is fairly unwieldy, with a lot of repetition – especially as we
 
 ```julia
 function linear(in, out)
-  W = param(randn(out, in))
-  b = param(randn(out))
+  W = randn(out, in)
+  b = randn(out)
   x -> W * x .+ b
 end
 
@@ -75,7 +131,7 @@ linear2 = linear(3, 2)
 
 model(x) = linear2(σ.(linear1(x)))
 
-model(x) # => 2-element vector
+model(rand(5)) # => 2-element vector
 ```
 
 Another (equivalent) way is to create a struct that explicitly represents the affine layer.
@@ -87,7 +143,7 @@ struct Affine
 end
 
 Affine(in::Integer, out::Integer) =
-  Affine(param(randn(out, in)), param(randn(out)))
+  Affine(randn(out, in), randn(out))
 
 # Overload call, so the object can be used as a function
 (m::Affine)(x) = m.W * x .+ m.b
@@ -118,7 +174,7 @@ using Flux
 
 layers = [Dense(10, 5, σ), Dense(5, 2), softmax]
 
-model(x) = foldl((x, m) -> m(x), x, layers)
+model(x) = foldl((x, m) -> m(x), layers, init = x)
 
 model(rand(10)) # => 2-element vector
 ```
@@ -151,3 +207,36 @@ m = Chain(x -> x^2, x -> x+1)
 
 m(5) # => 26
 ```
+
+## Layer helpers
+
+Flux provides a set of helpers for custom layers, which you can enable by calling
+
+```julia
+Flux.@functor Affine
+```
+
+This enables a useful extra set of functionality for our `Affine` layer, such as [collecting its parameters](../training/optimisers.md) or [moving it to the GPU](../gpu.md).
+
+For some more helpful tricks, including parameter freezing, please checkout the [advanced usage guide](advanced.md).
+
+## Utility functions
+
+Flux provides some utility functions to help you generate models in an automated fashion.
+
+`outdims` enables you to calculate the spatial output dimensions of layers like `Conv` when applied to input images of a given size.
+Currently limited to the following layers:
+- `Chain`
+- `Dense`
+- `Conv`
+- `Diagonal`
+- `Maxout`
+- `ConvTranspose`
+- `DepthwiseConv`
+- `CrossCor`
+- `MaxPool`
+- `MeanPool`
+
+```@docs
+Flux.outdims
+```

docs/src/models/layers.md

@@ -1,6 +1,92 @@
-## Model Layers
+## Basic Layers
+
+These core layers form the foundation of almost all neural networks.
 
 ```@docs
 Chain
 Dense
 ```
+
+## Convolution and Pooling Layers
+
+These layers are used to build convolutional neural networks (CNNs).
+
+```@docs
+Conv
+MaxPool
+GlobalMaxPool
+MeanPool
+GlobalMeanPool
+DepthwiseConv
+ConvTranspose
+CrossCor
+SamePad
+flatten
+Flux.Zeros
+Flux.convfilter
+Flux.depthwiseconvfilter
+```
+
+## Recurrent Layers
+
+Much like the core layers above, but can be used to process sequence data (as well as other kinds of structured data).
+
+```@docs
+RNN
+LSTM
+GRU
+Flux.Recur
+Flux.reset!
+```
+
+## Other General Purpose Layers
+These are marginally more obscure than the Basic Layers.
+But in contrast to the layers described in the other sections are not readily grouped around a particular purpose (e.g. CNNs or RNNs).
+
+```@docs
+Maxout
+SkipConnection
+```
+
+
+## Normalisation & Regularisation
+
+These layers don't affect the structure of the network but may improve training times or reduce overfitting.
+
+```@docs
+Flux.normalise
+BatchNorm
+Flux.dropout
+Dropout
+AlphaDropout
+LayerNorm
+InstanceNorm
+GroupNorm
+```
+
+### Testmode
+
+Many normalisation layers behave differently under training and inference (testing). By default, Flux will automatically determine when a layer evaluation is part of training or inference. Still, depending on your use case, it may be helpful to manually specify when these layers should be treated as being trained or not. For this, Flux provides `Flux.testmode!`. When called on a model (e.g. a layer or chain of layers), this function will place the model into the mode specified.
+
+```@docs
+Flux.testmode!
+trainmode!
+```
+
+## Cost Functions
+```@docs
+Flux.mae
+Flux.mse
+Flux.msle
+Flux.huber_loss
+Flux.crossentropy
+Flux.logitcrossentropy
+Flux.binarycrossentropy
+Flux.logitbinarycrossentropy
+Flux.kldivergence
+Flux.poisson
+Flux.hinge
+Flux.squared_hinge
+Flux.dice_coeff_loss
+Flux.tversky_loss
+```

docs/src/models/nnlib.md

@@ -0,0 +1,61 @@
+# NNlib
+
+Flux re-exports all of the functions exported by the [NNlib](https://github.com/FluxML/NNlib.jl) package.
+
+## Activation Functions
+
+Non-linearities that go between layers of your model. Note that, unless otherwise stated, activation functions operate on scalars. To apply them to an array you can call `σ.(xs)`, `relu.(xs)` and so on.
+
+```@docs
+NNlib.celu
+NNlib.elu
+NNlib.gelu
+NNlib.hardsigmoid
+NNlib.hardtanh
+NNlib.leakyrelu
+NNlib.lisht
+NNlib.logcosh
+NNlib.logsigmoid
+NNlib.mish
+NNlib.relu
+NNlib.relu6
+NNlib.rrelu
+NNlib.selu
+NNlib.sigmoid
+NNlib.softplus
+NNlib.softshrink
+NNlib.softsign
+NNlib.swish
+NNlib.tanhshrink
+NNlib.trelu
+```
+
+## Softmax
+
+```@docs
+NNlib.softmax
+NNlib.logsoftmax
+```
+
+## Pooling
+
+```@docs
+NNlib.maxpool
+NNlib.meanpool
+```
+
+## Convolution
+
+```@docs
+NNlib.conv
+NNlib.depthwiseconv
+```
+
+## Batched Operations
+
+```@docs
+NNlib.batched_mul
+NNlib.batched_mul!
+NNlib.batched_adjoint
+NNlib.batched_transpose
+```

docs/src/models/recurrence.md

@@ -77,7 +77,7 @@ If you use the `RNN(10, 5)` constructor – as opposed to `RNNCell` – you'll s
 
 ```julia
 julia> RNN(10, 5)
-Recur(RNNCell(Dense(15, 5)))
+Recur(RNNCell(10, 5, tanh))
 ```
 
 ## Sequences
@@ -101,14 +101,4 @@ m = Chain(LSTM(10, 15), Dense(15, 5))
 m.(seq)
 ```
 
-## Truncating Gradients
-
-By default, calculating the gradients in a recurrent layer involves the entire history. For example, if we call the model on 100 inputs, calling `back!` will calculate the gradient for those 100 calls. If we then calculate another 10 inputs we have to calculate 110 gradients – this accumulates and quickly becomes expensive.
-
-To avoid this we can *truncate* the gradient calculation, forgetting the history.
-
-```julia
-truncate!(m)
-```
-
-Calling `truncate!` wipes the slate clean, so we can call the model with more inputs without building up an expensive gradient computation.
+Finally, we can reset the hidden state of the cell back to its initial value using `reset!(m)`.

docs/src/models/regularisation.md

@@ -0,0 +1,70 @@
+# Regularisation
+
+Applying regularisation to model parameters is straightforward. We just need to
+apply an appropriate regulariser, such as `norm`, to each model parameter and
+add the result to the overall loss.
+
+For example, say we have a simple regression.
+
+```julia
+using Flux: crossentropy
+m = Dense(10, 5)
+loss(x, y) = crossentropy(softmax(m(x)), y)
+```
+
+We can regularise this by taking the (L2) norm of the parameters, `m.W` and `m.b`.
+
+```julia
+using LinearAlgebra
+
+penalty() = norm(m.W) + norm(m.b)
+loss(x, y) = crossentropy(softmax(m(x)), y) + penalty()
+```
+
+When working with layers, Flux provides the `params` function to grab all
+parameters at once. We can easily penalise everything with `sum(norm, params)`.
+
+```julia
+julia> params(m)
+2-element Array{Any,1}:
+ param([0.355408 0.533092; … 0.430459 0.171498])
+ param([0.0, 0.0, 0.0, 0.0, 0.0])
+
+julia> sum(norm, params(m))
+26.01749952921026
+```
+
+Here's a larger example with a multi-layer perceptron.
+
+```julia
+m = Chain(
+  Dense(28^2, 128, relu),
+  Dense(128, 32, relu),
+  Dense(32, 10), softmax)
+
+loss(x, y) = crossentropy(m(x), y) + sum(norm, params(m))
+
+loss(rand(28^2), rand(10))
+```
+
+One can also easily add per-layer regularisation via the `activations` function:
+
+```julia
+julia> using Flux: activations
+
+julia> c = Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
+Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
+
+julia> activations(c, rand(10))
+3-element Array{Any,1}:
+ Float32[0.84682214, 0.6704139, 0.42177814, 0.257832, 0.36255655]
+ Float32[0.1501253, 0.073269576]                                 
+ Float32[0.5192045, 0.48079553]                                  
+
+julia> sum(norm, ans)
+2.1166067f0
+```
+
+```@docs
+Flux.activations
+```

docs/src/performance.md

@@ -0,0 +1,76 @@
+# Performance Tips
+
+All the usual [Julia performance tips apply](https://docs.julialang.org/en/v1/manual/performance-tips/).
+As always [profiling your code](https://docs.julialang.org/en/v1/manual/profile/#Profiling-1) is generally a useful way of finding bottlenecks.
+Below follow some Flux specific tips/reminders.
+
+## Don't use more precision than you need
+
+Flux works great with all kinds of number types.
+But often you do not need to be working with say `Float64` (let alone `BigFloat`).
+Switching to `Float32` can give you a significant speed up,
+not because the operations are faster, but because the memory usage is halved.
+Which means allocations occur much faster.
+And you use less memory.
+
+
+## Preserve inputs' types
+
+Not only should your activation and loss functions be [type-stable](https://docs.julialang.org/en/v1/manual/performance-tips/#Write-%22type-stable%22-functions-1),
+they should also preserve the type of their inputs.
+
+A very artificial example using an activation function like
+
+```
+    my_tanh(x) = Float64(tanh(x))
+```
+
+will result in performance on `Float32` input orders of magnitude slower than the normal `tanh` would,
+because it results in having to use slow mixed type multiplication in the dense layers.
+Similar situations can occur in the loss function during backpropagation.
+
+Which means if you change your data say from `Float64` to `Float32` (which should give a speedup: see above),
+you will see a large slow-down.
+
+This can occur sneakily, because you can cause type-promotion by interacting with a numeric literals.
+E.g. the following will have run into the same problem as above:
+
+```
+    leaky_tanh(x) = 0.01*x + tanh(x)
+```
+
+While one could change the activation function (e.g. to use `0.01f0*x`), the idiomatic (and safe way)  to avoid type casts whenever inputs changes is to use `oftype`:
+```
+    leaky_tanh(x) = oftype(x/1, 0.01)*x + tanh(x)
+```
+
+
+## Evaluate batches as Matrices of features
+
+While it can sometimes be tempting to process your observations (feature vectors) one at a time
+e.g.
+```julia
+function loss_total(xs::AbstractVector{<:Vector}, ys::AbstractVector{<:Vector})
+    sum(zip(xs, ys)) do (x, y_target)
+        y_pred = model(x)  # evaluate the model
+        return loss(y_pred, y_target)
+    end
+end
+```
+
+It is much faster to concatenate them into a matrix,
+as this will hit BLAS matrix-matrix multiplication, which is much faster than the equivalent sequence of matrix-vector multiplications.
+The improvement is enough that it is worthwhile allocating new memory to store them contiguously.
+
+```julia
+x_batch = reduce(hcat, xs)
+y_batch = reduce(hcat, ys)
+...
+function loss_total(x_batch::Matrix, y_batch::Matrix)
+    y_preds = model(x_batch)
+    sum(loss.(y_preds, y_batch))
+end
+```
+
+When doing this kind of concatenation use `reduce(hcat, xs)` rather than `hcat(xs...)`.
+This will avoid the splatting penalty, and will hit the optimised `reduce` method.

docs/src/saving.md

@@ -0,0 +1,118 @@
+# Saving and Loading Models
+
+You may wish to save models so that they can be loaded and run in a later
+session. The easiest way to do this is via
+[BSON.jl](https://github.com/MikeInnes/BSON.jl).
+
+Save a model:
+
+```julia
+julia> using Flux
+
+julia> model = Chain(Dense(10,5,relu),Dense(5,2),softmax)
+Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)
+
+julia> using BSON: @save
+
+julia> @save "mymodel.bson" model
+```
+
+Load it again:
+
+```julia
+julia> using Flux
+
+julia> using BSON: @load
+
+julia> @load "mymodel.bson" model
+
+julia> model
+Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)
+```
+
+Models are just normal Julia structs, so it's fine to use any Julia storage
+format for this purpose. BSON.jl is particularly well supported and most likely
+to be forwards compatible (that is, models saved now will load in future
+versions of Flux).
+
+!!! note
+
+    If a saved model's weights are stored on the GPU, the model will not load
+    later on if there is no GPU support available. It's best to [move your model
+    to the CPU](gpu.md) with `cpu(model)` before saving it.
+
+## Saving Model Weights
+
+In some cases it may be useful to save only the model parameters themselves, and
+rebuild the model architecture in your code. You can use `params(model)` to get
+model parameters. You can also use `data.(params)` to remove tracking.
+
+```Julia
+julia> using Flux
+
+julia> model = Chain(Dense(10,5,relu),Dense(5,2),softmax)
+Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)
+
+julia> weights = params(model);
+
+julia> using BSON: @save
+
+julia> @save "mymodel.bson" weights
+```
+
+You can easily load parameters back into a model with `Flux.loadparams!`.
+
+```julia
+julia> using Flux
+
+julia> model = Chain(Dense(10,5,relu),Dense(5,2),softmax)
+Chain(Dense(10, 5, NNlib.relu), Dense(5, 2), NNlib.softmax)
+
+julia> using BSON: @load
+
+julia> @load "mymodel.bson" weights
+
+julia> Flux.loadparams!(model, weights)
+```
+
+The new `model` we created will now be identical to the one we saved parameters for.
+
+## Checkpointing
+
+In longer training runs it's a good idea to periodically save your model, so that you can resume if training is interrupted (for example, if there's a power cut). You can do this by saving the model in the [callback provided to `train!`](training/training.md).
+
+```julia
+using Flux: throttle
+using BSON: @save
+
+m = Chain(Dense(10,5,relu),Dense(5,2),softmax)
+
+evalcb = throttle(30) do
+  # Show loss
+  @save "model-checkpoint.bson" model
+end
+```
+
+This will update the `"model-checkpoint.bson"` file every thirty seconds.
+
+You can get more advanced by saving a series of models throughout training, for example
+
+```julia
+@save "model-$(now()).bson" model
+```
+
+will produce a series of models like `"model-2018-03-06T02:57:10.41.bson"`. You
+could also store the current test set loss, so that it's easy to (for example)
+revert to an older copy of the model if it starts to overfit.
+
+```julia
+@save "model-$(now()).bson" model loss = testloss()
+```
+
+You can even store optimiser state alongside the model, to resume training
+exactly where you left off.
+
+```julia
+opt = ADAM()
+@save "model-$(now()).bson" model opt
+```

docs/src/training/optimisers.md

@@ -3,52 +3,153 @@
 Consider a [simple linear regression](../models/basics.md). We create some dummy data, calculate a loss, and backpropagate to calculate gradients for the parameters `W` and `b`.
 
 ```julia
-W = param(rand(2, 5))
-b = param(rand(2))
+using Flux
 
-predict(x) = W*x .+ b
+W = rand(2, 5)
+b = rand(2)
+
+predict(x) = (W * x) .+ b
 loss(x, y) = sum((predict(x) .- y).^2)
 
 x, y = rand(5), rand(2) # Dummy data
 l = loss(x, y) # ~ 3
-back!(l)
+
+θ = Params([W, b])
+grads = gradient(() -> loss(x, y), θ)
 ```
 
 We want to update each parameter, using the gradient, in order to improve (reduce) the loss. Here's one way to do that:
 
 ```julia
-using Flux.Tracker: data, grad
+using Flux.Optimise: update!
 
-function update()
-  η = 0.1 # Learning Rate
-  for p in (W, b)
-    x, Δ = data(p), grad(p)
-    x .-= η .* Δ # Apply the update
-    Δ .= 0       # Clear the gradient
-  end
+η = 0.1 # Learning Rate
+for p in (W, b)
+  update!(p, -η * grads[p])
 end
 ```
 
-If we call `update`, the parameters `W` and `b` will change and our loss should go down.
-
-There are two pieces here: one is that we need a list of trainable parameters for the model (`[W, b]` in this case), and the other is the update step. In this case the update is simply gradient descent (`x .-= η .* Δ`), but we might choose to do something more advanced, like adding momentum.
-
-In this case, getting the variables is trivial, but you can imagine it'd be more of a pain with some complex stack of layers.
+Running this will alter the parameters `W` and `b` and our loss should go down. Flux provides a more general way to do optimiser updates like this.
 
 ```julia
-m = Chain(
-  Dense(10, 5, σ),
-  Dense(5, 2), softmax)
+opt = Descent(0.1) # Gradient descent with learning rate 0.1
+
+for p in (W, b)
+  update!(opt, p, grads[p])
+end
 ```
 
-Instead of having to write `[m[1].W, m[1].b, ...]`, Flux provides a params function `params(m)` that returns a list of all parameters in the model for you.
+An optimiser `update!` accepts a parameter and a gradient, and updates the parameter according to the chosen rule. We can also pass `opt` to our [training loop](training.md), which will update all parameters of the model in a loop. However, we can now easily replace `Descent` with a more advanced optimiser such as `ADAM`.
 
-For the update step, there's nothing whatsoever wrong with writing the loop above – it'll work just fine – but Flux provides various *optimisers* that make it more convenient.
+## Optimiser Reference
+
+All optimisers return an object that, when passed to `train!`, will update the parameters passed to it.
+
+```@docs
+Flux.Optimise.update!
+Descent
+Momentum
+Nesterov
+RMSProp
+ADAM
+RADAM
+AdaMax
+ADAGrad
+ADADelta
+AMSGrad
+NADAM
+ADAMW
+```
+
+## Optimiser Interface
+
+Flux's optimisers are built around a `struct` that holds all the optimiser parameters along with a definition of how to apply the update rule associated with it. We do this via the `apply!` function which takes the optimiser as the first argument followed by the parameter and its corresponding gradient.
+
+In this manner Flux also allows one to create custom optimisers to be used seamlessly. Let's work this with a simple example.
 
 ```julia
-opt = SGD([W, b], 0.1) # Gradient descent with learning rate 0.1
+mutable struct Momentum
+  eta
+  rho
+  velocity
+end
 
-opt()
+Momentum(eta::Real, rho::Real) = Momentum(eta, rho, IdDict())
 ```
 
-An optimiser takes a parameter list and returns a function that does the same thing as `update` above. We can pass either `opt` or `update` to our [training loop](training.md), which will then run the optimiser after every mini-batch of data.
+The `Momentum` type will act as our optimiser in this case. Notice that we have added all the parameters as fields, along with the velocity which we will use as our state dictionary. Each parameter in our models will get an entry in there. We can now define the rule applied when this optimiser is invoked.
+
+```julia
+function Flux.Optimise.apply!(o::Momentum, x, Δ)
+  η, ρ = o.eta, o.rho
+  v = get!(o.velocity, x, zero(x))::typeof(x)
+  @. v = ρ * v - η * Δ
+  @. Δ = -v
+end
+```
+
+This is the basic definition of a Momentum update rule given by:
+
+```math
+v = ρ * v - η * Δ
+w = w - v
+```
+
+The `apply!` defines the update rules for an optimiser `opt`, given the parameters and gradients. It returns the updated gradients. Here, every parameter `x` is retrieved from the running state `v` and subsequently updates the state of the optimiser.
+
+Flux internally calls on this function via the `update!` function. It shares the API with `apply!` but ensures that multiple parameters are handled gracefully.
+
+## Composing Optimisers
+
+Flux defines a special kind of optimiser simply called `Optimiser` which takes in arbitrary optimisers as input. Its behaviour is similar to the usual optimisers, but differs in that it acts by calling the optimisers listed in it sequentially. Each optimiser produces a modified gradient
+that will be fed into the next, and the resultant update will be applied to the parameter as usual. A classic use case is where adding decays is desirable. Flux defines some basic decays including `ExpDecay`, `InvDecay` etc.
+
+```julia
+opt = Optimiser(ExpDecay(0.001, 0.1, 1000, 1e-4), Descent())
+```
+
+Here we apply exponential decay to the `Descent` optimiser. The defaults of `ExpDecay` say that its learning rate will be decayed every 1000 steps.
+It is then applied like any optimiser.
+
+```julia
+w = randn(10, 10)
+w1 = randn(10,10)
+ps = Params([w, w1])
+
+loss(x) = Flux.mse(w * x, w1 * x)
+
+loss(rand(10)) # around 9
+
+for t = 1:10^5
+  θ = Params([w, w1])
+  θ̄ = gradient(() -> loss(rand(10)), θ)
+  Flux.Optimise.update!(opt, θ, θ̄)
+end
+
+loss(rand(10)) # around 0.9
+```
+
+In this manner it is possible to compose optimisers for some added flexibility.
+
+## Decays
+
+Similar to optimisers, Flux also defines some simple decays that can be used in conjunction with other optimisers, or standalone.
+
+```@docs
+ExpDecay
+InvDecay
+WeightDecay
+```
+
+## Gradient Clipping
+
+Gradient clipping is useful for training recurrent neural networks, which have a tendency to suffer from the exploding gradient problem. An example usage is
+
+```julia
+opt = Optimiser(ClipValue(1e-3), ADAM(1e-3))
+```
+
+```@docs
+ClipValue
+ClipNorm
+```

docs/src/training/training.md

@@ -1,36 +1,48 @@
 # Training
 
-To actually train a model we need three things:
+To actually train a model we need four things:
 
-* A *model loss function*, that evaluates how well a model is doing given some input data.
-* A collection of data points that will be provided to the loss function.
+* A *objective function*, that evaluates how well a model is doing given some input data.
+* The trainable parameters of the model.
+* A collection of data points that will be provided to the objective function.
 * An [optimiser](optimisers.md) that will update the model parameters appropriately.
 
-With these we can call `Flux.train!`:
+With these we can call `train!`:
 
-```julia
-Flux.train!(modelLoss, data, opt)
+```@docs
+Flux.Optimise.train!
 ```
 
 There are plenty of examples in the [model zoo](https://github.com/FluxML/model-zoo).
 
 ## Loss Functions
 
-The `loss` that we defined in [basics](../models/basics.md) is completely valid for training. We can also define a loss in terms of some model:
+The objective function must return a number representing how far the model is from its target – the *loss* of the model. The `loss` function that we defined in [basics](../models/basics.md) will work as an objective. We can also define an objective in terms of some model:
 
 ```julia
 m = Chain(
   Dense(784, 32, σ),
   Dense(32, 10), softmax)
 
-# Model loss function
 loss(x, y) = Flux.mse(m(x), y)
+ps = Flux.params(m)
 
 # later
-Flux.train!(loss, data, opt)
+Flux.train!(loss, ps, data, opt)
 ```
 
-The loss will almost always be defined in terms of some *cost function* that measures the distance of the prediction `m(x)` from the target `y`. Flux has several of these built in, like `mse` for mean squared error or `crossentropy` for cross entropy loss, but you can calculate it however you want.
+The objective will almost always be defined in terms of some *cost function* that measures the distance of the prediction `m(x)` from the target `y`. Flux has several of these built in, like `mse` for mean squared error or `crossentropy` for cross entropy loss, but you can calculate it however you want.
+For a list of all built-in loss functions, check out the [layer reference](../models/layers.md).
+
+At first glance it may seem strange that the model that we want to train is not part of the input arguments of `Flux.train!` too. However the target of the optimizer is not the model itself, but the objective function that represents the departure between modelled and observed data. In other words, the model is implicitly defined in the objective function, and there is no need to give it explicitly. Passing the objective function instead of the model and a cost function separately provides more flexibility, and the possibility of optimizing the calculations.
+
+## Model parameters
+
+The model to be trained must have a set of tracked parameters that are used to calculate the gradients of the objective function. In the [basics](../models/basics.md) section it is explained how to create models with such parameters. The second argument of the function `Flux.train!` must be an object containing those parameters, which can be obtained from a model `m` as `params(m)`.
+
+Such an object contains a reference to the model's parameters, not a copy, such that after their training, the model behaves according to their updated values.
+
+Handling all the parameters on a layer by layer basis is explained in the [Layer Helpers](../models/basics.md) section. Also, for freezing model parameters, see the [Advanced Usage Guide](../models/advanced.md).
 
 ## Datasets
 
@@ -47,7 +59,8 @@ data = [(x, y)]
 ```julia
 data = [(x, y), (x, y), (x, y)]
 # Or equivalently
-data = Iterators.repeated((x, y), 3)
+using IterTools: ncycle
+data = ncycle([(x, y)], 3)
 ```
 
 It's common to load the `x`s and `y`s separately. In this case you can use `zip`:
@@ -58,12 +71,40 @@ ys = [rand( 10), rand( 10), rand( 10)]
 data = zip(xs, ys)
 ```
 
+Training data can be conveniently  partitioned for mini-batch training using the [`Flux.Data.DataLoader`](@ref) type:
+
+```julia
+X = rand(28, 28, 60000)
+Y = rand(0:9, 60000)
+data = DataLoader(X, Y, batchsize=128) 
+```
+
+Note that, by default, `train!` only loops over the data once (a single "epoch").
+A convenient way to run multiple epochs from the REPL is provided by `@epochs`.
+
+```julia
+julia> using Flux: @epochs
+
+julia> @epochs 2 println("hello")
+INFO: Epoch 1
+hello
+INFO: Epoch 2
+hello
+
+julia> @epochs 2 Flux.train!(...)
+# Train for two epochs
+```
+
+```@docs
+Flux.@epochs
+```
+
 ## Callbacks
 
 `train!` takes an additional argument, `cb`, that's used for callbacks so that you can observe the training process. For example:
 
 ```julia
-train!(loss, data, opt, cb = () -> println("training"))
+train!(objective, ps, data, opt, cb = () -> println("training"))
 ```
 
 Callbacks are called for every batch of training data. You can slow this down using `Flux.throttle(f, timeout)` which prevents `f` from being called more than once every `timeout` seconds.
@@ -74,6 +115,41 @@ A more typical callback might look like this:
 test_x, test_y = # ... create single batch of test data ...
 evalcb() = @show(loss(test_x, test_y))
 
-Flux.train!(loss, data, opt,
+Flux.train!(objective, ps, data, opt,
             cb = throttle(evalcb, 5))
 ```
+
+Calling `Flux.stop()` in a callback will exit the training loop early.
+
+```julia
+cb = function ()
+  accuracy() > 0.9 && Flux.stop()
+end
+```
+
+## Custom Training loops
+
+The `Flux.train!` function can be very convenient, especially for simple problems.
+Its also very flexible with the use of callbacks.
+But for some problems its much cleaner to write your own custom training loop.
+An example follows that works similar to the default `Flux.train` but with no callbacks.
+You don't need callbacks if you just code the calls to your functions directly into the loop.
+E.g. in the places marked with comments.
+
+```julia
+function my_custom_train!(loss, ps, data, opt)
+  ps = Params(ps)
+  for d in data
+    gs = gradient(ps) do
+      training_loss = loss(d...)
+      # Insert whatever code you want here that needs Training loss, e.g. logging
+      return training_loss
+    end
+    # insert what ever code you want here that needs gradient
+    # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge
+    update!(opt, ps, gs)
+    # Here you might like to check validation set accuracy, and break out to do early stopping
+  end
+end
+```
+You could simplify this further, for example by hard-coding in the loss function.

docs/src/utilities.md

@@ -0,0 +1,49 @@
+# Utility Functions
+
+Flux contains some utility functions for working with data; these functions
+help create inputs for your models or batch your dataset.
+Other functions can be used to initialize your layers or to regularly execute
+callback functions.
+
+## Working with Data
+
+```@docs
+Flux.unsqueeze
+Flux.stack
+Flux.unstack
+Flux.chunk
+Flux.frequencies
+Flux.batch
+Flux.batchseq
+Base.rpad(v::AbstractVector, n::Integer, p)
+```
+
+## Layer Initialization
+
+These are primarily useful if you are planning to write your own layers.
+Flux initializes convolutional layers and recurrent cells with `glorot_uniform`
+by default.
+To change the default on an applicable layer, pass the desired function with the
+`init` keyword. For example:
+```jldoctest; setup = :(using Flux)
+julia> conv = Conv((3, 3), 1 => 8, relu; init=Flux.glorot_normal)
+Conv((3, 3), 1=>8, relu)
+```
+
+```@docs
+Flux.glorot_uniform
+Flux.glorot_normal
+```
+
+## Model Abstraction
+
+```@docs
+Flux.destructure
+```
+
+## Callback Helpers
+
+```@docs
+Flux.throttle
+Flux.stop
+```

paper/paper.bib

@@ -0,0 +1,50 @@
+@misc{Julia,
+  author = {Jeff Bezanson and Alan Edelman and Stefan Karpinski and Viral B. Shah},
+  title = {Julia: A Fresh Approach to Numerical Computing},
+  journal = {SIAM Review},
+  volume = {59},
+  year = {2017},
+  doi = {10.1137/141000671},
+  howpublished = {\url{julialang.org/publications/julia-fresh-approach-BEKS.pdf}}
+}
+
+@article{besard:2017,
+  author    = {Tim Besard and Christophe Foket and De Sutter, Bjorn},
+  title     = {Effective Extensible Programming: Unleashing {Julia} on {GPUs}},
+  journal   = {arXiv},
+  volume    = {abs/11712.03112},
+  year      = {2017},
+  url       = {https://arxiv.org/abs/1712.03112},
+}
+
+@online{MLPL,
+  author = {Mike Innes and others},
+  title = {On Machine Learning and Programming Languages},
+  year = 2017,
+  url = {https://julialang.org/blog/2017/12/ml&pl},
+  urldate = {2018-02-16}
+}
+
+@online{CuArrays,
+  author = {Mike Innes and others},
+  title = {Generic GPU Kernels},
+  year = 2017,
+  url = {https://mikeinnes.github.io/2017/08/24/cudanative.html},
+  urldate = {2018-02-16}
+}
+
+@online{Zoo,
+  author = {Mike Innes and others},
+  title = {Flux Model Zoo},
+  year = 2018,
+  url = {https://github.com/FluxML/model-zoo/},
+  urldate = {2018-02-16}
+}
+
+@online{Minibatch,
+  author = {James Bradbury},
+  title = {Minibatch.jl},
+  year = 2018,
+  url = {https://github.com/jekbradbury/Minibatch.jl},
+  urldate = {2018-02-16}
+}

paper/paper.md

@@ -0,0 +1,31 @@
+---
+title: 'Flux: Elegant machine learning with Julia'
+tags:
+  - deep learning
+  - machine learning
+  - natural language processing
+  - computer vision
+  - reinforcement learning
+  - robotics
+  - automatic differentiation
+  - compiler
+authors:
+  - name: Mike Innes
+    orcid: 0000-0003-0788-0242
+    affiliation: 1
+affiliations:
+  - name: Julia Computing
+    index: 1
+date: 16 February 2018
+bibliography: paper.bib
+---
+
+# Summary
+
+Flux is library for machine learning (ML), written using the numerical computing language Julia [@Julia]. The package allows models to be written using Julia's simple mathematical syntax, and applies automatic differentiation (AD) to seamlessly calculate derivatives and train the model. Meanwhile, it makes heavy use of Julia's language and compiler features to carry out code analysis and make optimisations. For example, Julia's GPU compilation support [@besard:2017] can be used to JIT-compile custom GPU kernels for model layers [@CuArrays].
+
+The machine learning community has traditionally been divided between "static" and "dynamic" frameworks that are easy to optimise and easy to use, respectively [@MLPL]. Flux blurs the line between these two approaches, combining a highly intuitive programming model with the compiler techniques needed by ML. This enables research into advanced compiler transforms such as batching [@Minibatch] without changing any user code.
+
+Flux has been used heavily for natural language processing, but can also support state-of-the-art research models in areas like computer vision, reinforcement learning and robotics. Many examples of such models can be found in the model zoo [@Zoo].
+
+# References

src/Flux.jl

@@ -1,31 +1,57 @@
-__precompile__()
-
 module Flux
 
 # Zero Flux Given
 
-using Juno, Requires
-using Lazy: @forward
+using Base: tail
+using Statistics, Random, LinearAlgebra
+using Zygote, MacroTools, Juno, Reexport
+using MacroTools: @forward
+@reexport using NNlib
+using Zygote: Params, @adjoint, gradient, pullback, @nograd
 
-export Chain, Dense, RNN, LSTM,
-  SGD, param, params, mapleaves
+export gradient
 
-using NNlib
-export σ, relu, softmax
-
-include("tracker/Tracker.jl")
-using .Tracker
+export Chain, Dense, Maxout, RNN, LSTM, GRU, SamePad, Conv, CrossCor, ConvTranspose,
+       GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool, flatten,
+       DepthwiseConv, Dropout, AlphaDropout, LayerNorm, BatchNorm, InstanceNorm, GroupNorm,
+       SkipConnection, params, fmap, cpu, gpu, f32, f64, testmode!, trainmode!
 
 include("optimise/Optimise.jl")
 using .Optimise
+using .Optimise: @epochs
+export Descent, ADAM, Momentum, Nesterov, RMSProp,
+  ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM,
+  ADAMW, RADAM, InvDecay, ExpDecay, WeightDecay,
+  ClipValue, ClipNorm
+
+
+using CuArrays
+const use_cuda = Ref(false)
 
 include("utils.jl")
+include("zeros.jl")
 include("onehot.jl")
-include("tree.jl")
+include("functor.jl")
 
-include("layers/softmax.jl")
 include("layers/stateless.jl")
 include("layers/basic.jl")
+include("layers/conv.jl")
 include("layers/recurrent.jl")
+include("layers/normalise.jl")
+
+include("data/Data.jl")
+
+include("deprecations.jl")
+
+include("cuda/cuda.jl")
+
+function __init__()
+  use_cuda[] = CuArrays.functional() # Can be overridden after load with `Flux.use_cuda[] = false`
+  if CuArrays.functional()
+    if !CuArrays.has_cudnn()
+      @warn "CuArrays.jl found cuda, but did not find libcudnn. Some functionality will not be available."
+    end
+  end
+end
 
 end # module

src/cuda/cuda.jl

@@ -0,0 +1,9 @@
+module CUDA
+
+using ..CuArrays
+
+using CuArrays: CUDNN
+include("curnn.jl")
+include("cudnn.jl")
+
+end

src/cuda/cudnn.jl

@@ -0,0 +1,8 @@
+import ..Flux: data
+import CuArrays.CUDNN: batchnorm, ∇batchnorm
+
+(BN::Flux.BatchNorm)(x::Union{CuArray{T,2},CuArray{T,4},CuArray{T,5}}, cache = nothing) where T<:Union{Float32, Float64} =
+  BN.λ.(batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ², BN.momentum; cache = cache, alpha = 1, beta = 0, eps = BN.ϵ, training = Flux.istraining()))
+
+@adjoint batchnorm(g, b, x, running_mean, running_var, momentum; kw...) =
+  batchnorm(g, b, x, running_mean, running_var, momentum; kw...), Δ -> (∇batchnorm(g, b, x, Δ, running_mean, running_var, momentum; kw...)..., nothing, nothing, nothing)

src/cuda/curnn.jl

@@ -0,0 +1,90 @@
+import ..Flux: Flux, relu
+using CuArrays.CUDAnative
+
+CuRNN{T} = Flux.RNNCell{<:Union{typeof(tanh),typeof(relu)},<:CuArray{T,2},<:CuArray{T,1}}
+CuGRU{T} = Flux.GRUCell{<:CuArray{T,2},<:CuArray{T,1}}
+CuLSTM{T} = Flux.LSTMCell{<:CuArray{T,2},<:CuArray{T,1}}
+CuRNNs{T} = Union{CuRNN{T},CuGRU{T},CuLSTM{T}}
+
+function CUDNN.RNNDesc(m::CuRNNs{T}) where T
+  h, i = length(m.h), size(m.Wi, 2)
+  mode = m isa CuRNN ?
+    (m.σ == tanh ? CUDNN.CUDNN_RNN_TANH : CUDNN.CUDNN_RNN_RELU) :
+    m isa CuGRU ? CUDNN.CUDNN_GRU : CUDNN.CUDNN_LSTM
+  r = CUDNN.RNNDesc{T}(mode, i, h)
+  return r
+end
+
+const descs = WeakKeyDict()
+
+function desc(rnn)
+  d = haskey(descs, rnn) ? descs[rnn] : (descs[rnn] = CUDNN.RNNDesc(rnn))
+  CUDNN.setweights!(d, rnn.Wi, rnn.Wh, rnn.b)
+  return d
+end
+
+import Zygote
+using Zygote: @adjoint
+
+function (m::CuRNN{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
+  y, h′ = CUDNN.forward(desc(m), x, h)
+  return h′, y
+end
+
+function (m::CuGRU{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
+  y, h′ = CUDNN.forward(desc(m), x, h)
+  return h′, y
+end
+
+function (m::CuLSTM{T})(h::NTuple{2,CuArray{T}}, x::CuArray{T}) where T <: Union{Float32,Float64}
+  y, h′, c′ = CUDNN.forward(desc(m), x, h[1], h[2])
+  return (h′, c′), y
+end
+
+(m::CuRNN{T})(h::CuArray{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
+(m::CuGRU{T})(h::CuArray{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
+(m::CuLSTM{T})(h::NTuple{2,CuArray{T}}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
+
+trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))
+
+unbroadcast(x::AbstractArray, Δ) =
+  size(x) == size(Δ) ? Δ :
+  length(x) == length(Δ) ? trim(x, Δ) :
+    trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ)))))
+
+coerce_cuda(x::Union{CuArray,Nothing}) = x
+coerce_cuda(x::Tuple) = coerce_cuda.(x)
+
+coerce_cuda(x::AbstractArray) = x .+ CuArrays.fill(0)
+
+function struct_grad!(cx::Zygote.Context, x, x̄)
+  for f in fieldnames(typeof(x))
+    Zygote.accum_param(cx, getfield(x, f), getfield(x̄, f))
+  end
+  dx = Zygote.grad_mut(cx, x)
+  dx[] = Zygote.accum(dx[], x̄)
+  return dx
+end
+
+for RNN in (CuRNN, CuGRU)
+  @eval @adjoint function (m::$RNN{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
+    (y, ho), back = CUDNN.pullback(desc(m), x, h)
+    (ho, y), function (Δ)
+      dho, dy = coerce_cuda(Δ) # Support FillArrays etc.
+      m̄ = back(dy, dho)
+      dm = struct_grad!(__context__, m, (σ=nothing,Wi=transpose(m̄.Wi),Wh=transpose(m̄.Wh),b=m̄.b,h=nothing))
+      (dm, unbroadcast(h, m̄.h), m̄.x)
+    end
+  end
+end
+
+@adjoint function (m::CuLSTM)((h, c)::Tuple{CuArray{T},CuArray{T}}, x::CuArray{T}) where T <: Union{Float32,Float64}
+  (y, ho, co), back = CUDNN.pullback(desc(m), x, h, c)
+  ((ho, co), y), function (Δ)
+    dhc, dy = coerce_cuda(Δ) # Support FillArrays etc.
+    dho, dco = dhc === nothing ? (nothing, nothing) : dhc
+    m̄ = back(dy, dho, dco)
+    dm = struct_grad!(__context__, m, (σ=nothing,Wi=transpose(m̄.Wi),Wh=transpose(m̄.Wh),b=m̄.b,h=nothing,c=nothing))
+    (dm, (unbroadcast(h, m̄.h), unbroadcast(c, m̄.c)), m̄.x)
+  end
+end

src/data/Data.jl

@@ -0,0 +1,56 @@
+module Data
+
+import ..Flux
+import SHA
+
+using Random: shuffle!
+using Base: @propagate_inbounds
+
+export CMUDict, cmudict
+
+deps(path...) = joinpath(@__DIR__, "..", "..", "deps", path...)
+
+function download_and_verify(url, path, hash)
+    tmppath = tempname()
+    download(url, tmppath)
+    hash_download = open(tmppath) do f
+        bytes2hex(SHA.sha256(f))
+    end
+    if hash_download !== hash
+        msg  = "Hash Mismatch!\n"
+        msg *= "  Expected sha256:   $hash\n"
+        msg *= "  Calculated sha256: $hash_download"
+        error(msg)
+    end
+    mv(tmppath, path; force=true)
+end
+
+function __init__()
+  mkpath(deps())
+end
+
+include("dataloader.jl")
+export DataLoader
+
+include("mnist.jl")
+export MNIST
+
+include("fashion-mnist.jl")
+export FashionMNIST
+
+include("cmudict.jl")
+using .CMUDict
+
+include("tree.jl")
+include("sentiment.jl")
+using .Sentiment
+
+include("iris.jl")
+export Iris
+
+include("housing.jl")
+export Housing
+
+@deprecate DataLoader(x...; kws...) DataLoader(x; kws...)
+
+end

src/data/cmudict.jl

@@ -0,0 +1,76 @@
+module CMUDict
+
+export cmudict
+
+using ..Data: deps, download_and_verify
+
+const version = "0.7b"
+const cache_prefix = "https://cache.julialang.org"
+
+function load()
+  suffixes_and_hashes = [(""       , "209a8b4cd265013e96f4658632a9878103b0c5abf62b50d4ef3ae1be226b29e4"),
+                        (".phones" , "ffb588a5e55684723582c7256e1d2f9fadb130011392d9e59237c76e34c2cfd6"),
+                        (".symbols", "408ccaae803641c6d7b626b6299949320c2dbca96b2220fd3fb17887b023b027")]
+  if isdir(deps("cmudict"))
+    if all(isfile(deps("cmudict", "cmudict$x")) for (x, _) in suffixes_and_hashes)
+      return
+    end
+  end
+  @info "Downloading CMUDict dataset"
+  mkpath(deps("cmudict"))
+  for (x, hash) in suffixes_and_hashes
+    download_and_verify("$cache_prefix/https://svn.code.sf.net/p/cmusphinx/code/trunk/cmudict/cmudict-$version$x",
+             deps("cmudict", "cmudict$x"), hash)
+  end
+end
+
+"""
+    phones()
+
+Return a `Vector` containing the phones used in the CMU Pronouncing Dictionary.
+"""
+function phones()
+  load()
+  Symbol.(first.(split.(split(read(deps("cmudict", "cmudict.phones"),String),
+                        "\n", keepempty = false), "\t")))
+end
+
+"""
+    symbols()
+
+Return a `Vector` containing the symbols used in the CMU Pronouncing Dictionary.
+A symbol is a phone with optional auxiliary symbols, indicating for example the
+amount of stress on the phone.
+"""
+function symbols()
+  load()
+  Symbol.(split(read(deps("cmudict", "cmudict.symbols"),String),
+                "\n", keepempty = false))
+end
+
+"""
+    rawdict()
+
+Return the unfiltered CMU Pronouncing Dictionary.
+"""
+function rawdict()
+  load()
+  Dict(String(xs[1]) => Symbol.(xs[2:end]) for xs in
+       filter(!isempty, split.(split(read(deps("cmudict", "cmudict"),String), "\n"))))
+end
+
+validword(s) = isascii(s) && occursin(r"^[\w\-\.]+$", s)
+
+"""
+    cmudict()
+
+Return a filtered CMU Pronouncing Dictionary.
+
+It is filtered so each word contains only ASCII characters and a combination of
+word characters (as determined by the regex engine using `\\w`), '-' and '.'.
+"""
+cmudict() = filter(p -> validword(p.first), rawdict())
+
+alphabet() = ['A':'Z'..., '0':'9'..., '_', '-', '.']
+
+end

src/data/dataloader.jl

@@ -0,0 +1,110 @@
+# Adapted from Knet's src/data.jl (author: Deniz Yuret)
+
+struct DataLoader{D}
+    data::D
+    batchsize::Int
+    nobs::Int
+    partial::Bool
+    imax::Int
+    indices::Vector{Int}
+    shuffle::Bool
+end
+
+"""
+    DataLoader(data; batchsize=1, shuffle=false, partial=true)
+
+An object that iterates over mini-batches of `data`, each mini-batch containing `batchsize` observations
+(except possibly the last one). 
+
+Takes as input a single data tensor, or a tuple (or a named tuple) of tensors.
+The last dimension in each tensor is considered to be the observation dimension.
+
+If `shuffle=true`, shuffles the observations each time iterations are re-started.
+If `partial=false`, drops the last mini-batch if it is smaller than the batchsize.
+
+The original data is preserved in the `data` field of the DataLoader. 
+
+Usage example:
+
+    Xtrain = rand(10, 100)
+    train_loader = DataLoader(Xtrain, batchsize=2) 
+    # iterate over 50 mini-batches of size 2
+    for x in train_loader
+        @assert size(x) == (10, 2)
+        ...
+    end
+
+    train_loader.data   # original dataset
+
+    # similar, but yielding tuples
+    train_loader = DataLoader((Xtrain,), batchsize=2) 
+    for (x,) in train_loader
+        @assert size(x) == (10, 2)
+        ...
+    end
+
+    Xtrain = rand(10, 100)
+    Ytrain = rand(100)
+    train_loader = DataLoader((Xtrain, Ytrain), batchsize=2, shuffle=true) 
+    for epoch in 1:100
+        for (x, y) in train_loader
+            @assert size(x) == (10, 2)
+            @assert size(y) == (2,)
+            ...
+        end
+    end
+
+    # train for 10 epochs
+    using IterTools: ncycle 
+    Flux.train!(loss, ps, ncycle(train_loader, 10), opt)
+
+    # can use NamedTuple to name tensors
+    train_loader = DataLoader((images=Xtrain, labels=Ytrain), batchsize=2, shuffle=true)
+    for datum in train_loader
+        @assert size(datum.images) == (10, 2)
+        @assert size(datum.labels) == (2,)
+    end
+"""
+function DataLoader(data; batchsize=1, shuffle=false, partial=true)
+    batchsize > 0 || throw(ArgumentError("Need positive batchsize"))
+    
+    n = _nobs(data) 
+    if n < batchsize
+        @warn "Number of observations less than batchsize, decreasing the batchsize to $n"
+        batchsize = n
+    end
+    imax = partial ? n : n - batchsize + 1
+    DataLoader(data, batchsize, n, partial, imax, [1:n;], shuffle)
+end
+
+@propagate_inbounds function Base.iterate(d::DataLoader, i=0)     # returns data in d.indices[i+1:i+batchsize]
+    i >= d.imax && return nothing
+    if d.shuffle && i == 0
+        shuffle!(d.indices)
+    end
+    nexti = min(i + d.batchsize, d.nobs)
+    ids = d.indices[i+1:nexti]
+    batch = _getobs(d.data, ids)
+    return (batch, nexti)
+end
+
+function Base.length(d::DataLoader)
+    n = d.nobs / d.batchsize
+    d.partial ? ceil(Int,n) : floor(Int,n)
+end
+
+_nobs(data::AbstractArray) = size(data)[end]
+
+function _nobs(data::Union{Tuple, NamedTuple})
+    length(data) > 0 || throw(ArgumentError("Need at least one data input"))
+    n = _nobs(data[1])
+    if !all(x -> _nobs(x) == n, Base.tail(data))
+        throw(DimensionMismatch("All data should contain same number of observations"))
+    end
+    return n
+end
+
+_getobs(data::AbstractArray, i) = data[ntuple(i -> Colon(), Val(ndims(data) - 1))..., i]
+_getobs(data::Union{Tuple, NamedTuple}, i) = map(Base.Fix2(_getobs, i), data)
+
+Base.eltype(::DataLoader{D}) where D = D

src/data/fashion-mnist.jl

@@ -0,0 +1,66 @@
+module FashionMNIST
+
+using ..MNIST: gzopen, imageheader, rawimage, labelheader, rawlabel
+using ..Data: download_and_verify
+
+const dir = joinpath(@__DIR__, "../../deps/fashion-mnist")
+
+function load()
+  mkpath(dir)
+  cd(dir) do
+    for (file, hash) in [("train-images-idx3-ubyte", "3aede38d61863908ad78613f6a32ed271626dd12800ba2636569512369268a84"),
+                         ("train-labels-idx1-ubyte", "a04f17134ac03560a47e3764e11b92fc97de4d1bfaf8ba1a3aa29af54cc90845"),
+                         ("t10k-images-idx3-ubyte" , "346e55b948d973a97e58d2351dde16a484bd415d4595297633bb08f03db6a073"),
+                         ("t10k-labels-idx1-ubyte" , "67da17c76eaffca5446c3361aaab5c3cd6d1c2608764d35dfb1850b086bf8dd5")]
+      isfile(file) && continue
+      @info "Downloading Fashion-MNIST dataset"
+      download_and_verify("http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/$file.gz", "$file.gz", hash)
+      open(file, "w") do io
+        write(io, gzopen(read, "$file.gz"))
+      end
+    end
+  end
+end
+
+const TRAINIMAGES = joinpath(dir, "train-images-idx3-ubyte")
+const TRAINLABELS = joinpath(dir, "train-labels-idx1-ubyte")
+const TESTIMAGES = joinpath(dir, "t10k-images-idx3-ubyte")
+const TESTLABELS = joinpath(dir, "t10k-labels-idx1-ubyte")
+
+"""
+    images()
+    images(:test)
+
+Load the Fashion-MNIST images.
+
+Each image is a 28×28 array of `Gray` colour values
+(see [Colors.jl](https://github.com/JuliaGraphics/Colors.jl)).
+
+Return the 60,000 training images by default; pass `:test` to retrieve the
+10,000 test images.
+"""
+function images(set = :train)
+  load()
+  io = IOBuffer(read(set == :train ? TRAINIMAGES : TESTIMAGES))
+  _, N, nrows, ncols = imageheader(io)
+  [rawimage(io) for _ in 1:N]
+end
+
+"""
+    labels()
+    labels(:test)
+
+Load the labels corresponding to each of the images returned from [`images()`](@ref).
+Each label is a number from 0-9.
+
+Return the 60,000 training labels by default; pass `:test` to retrieve the
+10,000 test labels.
+"""
+function labels(set = :train)
+  load()
+  io = IOBuffer(read(set == :train ? TRAINLABELS : TESTLABELS))
+  _, N = labelheader(io)
+  [rawlabel(io) for _ = 1:N]
+end
+
+end

src/data/housing.jl

@@ -0,0 +1,136 @@
+"""
+1. Title: Boston Housing Data
+
+2. Sources:
+   (a) Origin:  This dataset was taken from the StatLib library which is
+                maintained at Carnegie Mellon University.
+   (b) Creator:  Harrison, D. and Rubinfeld, D.L. 'Hedonic prices and the 
+                 demand for clean air', J. Environ. Economics & Management,
+                 vol.5, 81-102, 1978.
+   (c) Date: July 7, 1993
+
+3. Number of Instances: 506
+
+4. Number of Attributes: 13 continuous attributes (including "class"
+                            attribute "MEDV"), 1 binary-valued attribute.
+                            
+5. Attribute Information:
+   
+       1. CRIM      per capita crime rate by town
+       2. ZN        proportion of residential land zoned for lots over 
+                    25,000 sq.ft.
+       3. INDUS     proportion of non-retail business acres per town
+       4. CHAS      Charles River dummy variable (= 1 if tract bounds 
+                    river; 0 otherwise)
+       5. NOX       nitric oxides concentration (parts per 10 million)
+       6. RM        average number of rooms per dwelling
+       7. AGE       proportion of owner-occupied units built prior to 1940
+       8. DIS       weighted distances to five Boston employment centres
+       9. RAD       index of accessibility to radial highways
+       10. TAX      full-value property-tax rate per 10,000 dollars
+       11. PTRATIO  pupil-teacher ratio by town
+       12. B        1000(Bk - 0.63)^2 where Bk is the proportion of blacks 
+                    by town
+       13. LSTAT    % lower status of the population
+       14. MEDV     Median value of owner-occupied homes in 1000's of dollars   
+
+       Downloaded From: https://archive.ics.uci.edu/ml/machine-learning-databases/housing/housing.data
+
+"""
+module Housing
+
+using DelimitedFiles
+using ..Data: deps, download_and_verify
+
+#Uncomment if package exists
+#const cache_prefix = "https://cache.julialang.org/"
+const cache_prefix = ""
+
+function load()
+    isfile(deps("housing.data")) && return
+    
+    @info "Downloading the Boston housing Dataset"
+    download_and_verify("$(cache_prefix)http://archive.ics.uci.edu/ml/machine-learning-databases/housing/housing.data",
+                        deps("housing.data"),
+                        "baadf72995725d76efe787b664e1f083388c79ba21ef9a7990d87f774184735a")
+    
+    #@info "Download complete. Working on the files"
+    path = deps()
+    isfile(deps("housing.data")) && touch(joinpath(path, "tempfile.data"))
+    open(joinpath(path, "tempfile.data"), "a") do fout
+        open(deps("housing.data"), "r") do fin
+            for line in eachline(fin)
+                line = replace(lstrip(line), r" +" => s",")
+                println(fout, line)
+            end
+        end
+    end
+    mv(joinpath(path, "tempfile.data"), deps("housing.data"), force=true)
+end
+
+"""
+Gets the targets for the Boston housing dataset, a 506 element array listing the targets for each example
+
+```jldoctest
+julia> using Flux
+
+julia> target = Flux.Data.Housing.targets()
+
+julia> summary(target)
+506×1 Array{Float64,2}
+
+julia> target[1]
+24.0
+
+"""
+function targets()
+    load()
+    housing = readdlm(deps("housing.data"), ',')
+    reshape(Vector{Float64}(housing[1:end,end]), (506, 1))           
+end
+
+
+"""
+Gets the names of the features provided in the dataset
+
+"""
+function feature_names()
+    ["crim","zn","indus","chas","nox","rm","age","dis","rad","tax","ptratio","b","lstat"]
+end
+
+
+"""
+Gets the features of the Boston Housing Dataset. This is a 506x13 Matrix of Float64 datatypes.
+The values are in the order ["crim","zn","indus","chas","nox","rm","age","dis","rad","tax","ptratio","b","lstat"].
+It has 506 examples.
+
+```jldoctest
+julia> using Flux
+
+julia> features = Flux.Data.Housing.features()
+
+julia> summary(features)
+506×13 Array{Float64,2}
+
+julia> features[1, :]
+13-element Array{Float64,1}:
+0.00632
+18.0    
+2.31   
+0.0    
+0.538  
+   ⋮      
+296.0    
+15.3    
+396.9    
+4.98   
+
+"""
+function features()
+    load()
+    housing = readdlm(deps("housing.data"), ',')
+    Matrix{Float64}(housing[1:end, 1:13])    
+end
+
+
+end

src/data/iris.jl

@@ -0,0 +1,78 @@
+"""
+Fisher's classic iris dataset.
+
+Measurements from 3 different species of iris: setosa, versicolor and
+virginica. There are 50 examples of each species.
+
+There are 4 measurements for each example: sepal length, sepal width,
+petal length and petal width. The measurements are in centimeters.
+
+The module retrieves the data from the [UCI Machine Learning Repository](https://archive.ics.uci.edu/ml/datasets/iris).
+"""
+module Iris
+
+using DelimitedFiles
+using ..Data: deps, download_and_verify
+
+# Uncomment if the iris.data file is cached to cache.julialang.org.
+const cache_prefix = "https://cache.julialang.org/"
+
+function load()
+    isfile(deps("iris.data")) && return
+
+    @info "Downloading iris dataset."
+    download_and_verify("$(cache_prefix)https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data",
+                        deps("iris.data"),
+                        "6f608b71a7317216319b4d27b4d9bc84e6abd734eda7872b71a458569e2656c0")
+end
+
+"""
+    labels()
+
+Get the labels of the iris dataset, a 150 element array of strings listing the
+species of each example.
+
+```jldoctest; setup = :(Flux.Data.Iris.load())
+julia> labels = Flux.Data.Iris.labels();
+
+julia> summary(labels)
+"150-element Array{String,1}"
+
+julia> labels[1]
+"Iris-setosa"
+```
+"""
+function labels()
+    load()
+    iris = readdlm(deps("iris.data"), ',')
+    Vector{String}(iris[1:end, end])
+end
+
+"""
+    features()
+
+Get the features of the iris dataset. This is a 4x150 matrix of Float64
+elements. It has a row for each feature (sepal length, sepal width,
+petal length, petal width) and a column for each example.
+
+```jldoctest; setup = :(Flux.Data.Iris.load())
+julia> features = Flux.Data.Iris.features();
+
+julia> summary(features)
+"4×150 Array{Float64,2}"
+
+julia> features[:, 1]
+4-element Array{Float64,1}:
+ 5.1
+ 3.5
+ 1.4
+ 0.2
+```
+"""
+function features()
+    load()
+    iris = readdlm(deps("iris.data"), ',')
+    Matrix{Float64}(iris[1:end, 1:4]')
+end
+
+end

src/data/mnist.jl

@@ -0,0 +1,116 @@
+module MNIST
+
+using CodecZlib, Colors
+using ..Data: download_and_verify
+
+const Gray = Colors.Gray{Colors.N0f8}
+
+const dir = joinpath(@__DIR__, "../../deps/mnist")
+
+function gzopen(f, file)
+  open(file) do io
+    f(GzipDecompressorStream(io))
+  end
+end
+
+function load()
+  mkpath(dir)
+  cd(dir) do
+    for (file, hash) in [("train-images-idx3-ubyte", "440fcabf73cc546fa21475e81ea370265605f56be210a4024d2ca8f203523609"),
+                         ("train-labels-idx1-ubyte", "3552534a0a558bbed6aed32b30c495cca23d567ec52cac8be1a0730e8010255c"),
+                         ("t10k-images-idx3-ubyte" , "8d422c7b0a1c1c79245a5bcf07fe86e33eeafee792b84584aec276f5a2dbc4e6"),
+                         ("t10k-labels-idx1-ubyte" , "f7ae60f92e00ec6debd23a6088c31dbd2371eca3ffa0defaefb259924204aec6")]
+      isfile(file) && continue
+      @info "Downloading MNIST dataset"
+      download_and_verify("https://cache.julialang.org/http://yann.lecun.com/exdb/mnist/$file.gz", "$file.gz", hash)
+      open(file, "w") do io
+        write(io, gzopen(read, "$file.gz"))
+      end
+    end
+  end
+end
+
+const IMAGEOFFSET = 16
+const LABELOFFSET = 8
+
+const NROWS = 28
+const NCOLS = 28
+
+const TRAINIMAGES = joinpath(dir, "train-images-idx3-ubyte")
+const TRAINLABELS = joinpath(dir, "train-labels-idx1-ubyte")
+const TESTIMAGES = joinpath(dir, "t10k-images-idx3-ubyte")
+const TESTLABELS = joinpath(dir, "t10k-labels-idx1-ubyte")
+
+function imageheader(io::IO)
+  magic_number = bswap(read(io, UInt32))
+  total_items = bswap(read(io, UInt32))
+  nrows = bswap(read(io, UInt32))
+  ncols = bswap(read(io, UInt32))
+  return magic_number, Int(total_items), Int(nrows), Int(ncols)
+end
+
+function labelheader(io::IO)
+  magic_number = bswap(read(io, UInt32))
+  total_items = bswap(read(io, UInt32))
+  return magic_number, Int(total_items)
+end
+
+function rawimage(io::IO)
+  img = Array{Gray}(undef, NCOLS, NROWS)
+  for i in 1:NCOLS, j in 1:NROWS
+    img[i, j] = reinterpret(Colors.N0f8, read(io, UInt8))
+  end
+  return img
+end
+
+function rawimage(io::IO, index::Integer)
+  seek(io, IMAGEOFFSET + NROWS * NCOLS * (index - 1))
+  return rawimage(io)
+end
+
+rawlabel(io::IO) = Int(read(io, UInt8))
+
+function rawlabel(io::IO, index::Integer)
+  seek(io, LABELOFFSET + (index - 1))
+  return rawlabel(io)
+end
+
+getfeatures(io::IO, index::Integer) = vec(getimage(io, index))
+
+"""
+    images()
+    images(:test)
+
+Load the MNIST images.
+
+Each image is a 28×28 array of `Gray` colour values
+(see [Colors.jl](https://github.com/JuliaGraphics/Colors.jl)).
+
+Return the 60,000 training images by default; pass `:test` to retrieve the
+10,000 test images.
+"""
+function images(set = :train)
+  load()
+  io = IOBuffer(read(set == :train ? TRAINIMAGES : TESTIMAGES))
+  _, N, nrows, ncols = imageheader(io)
+  [rawimage(io) for _ in 1:N]
+end
+
+"""
+    labels()
+    labels(:test)
+
+Load the labels corresponding to each of the images returned from [`images()`](@ref).
+Each label is a number from 0-9.
+
+Return the 60,000 training labels by default; pass `:test` to retrieve the
+10,000 test labels.
+"""
+function labels(set = :train)
+  load()
+  io = IOBuffer(read(set == :train ? TRAINLABELS : TESTLABELS))
+  _, N = labelheader(io)
+  [rawlabel(io) for _ = 1:N]
+end
+
+end # module

src/data/sentiment.jl

@@ -0,0 +1,67 @@
+"Stanford Sentiment Treebank dataset."
+module Sentiment
+
+using ZipFile
+using ..Data: deps, download_and_verify
+
+function load()
+  isfile(deps("sentiment.zip")) && return
+  @info "Downloading sentiment treebank dataset"
+  download_and_verify("https://cache.julialang.org/https://nlp.stanford.edu/sentiment/trainDevTestTrees_PTB.zip",
+           deps("sentiment.zip"), "5c613a4f673fc74097d523a2c83f38e0cc462984d847b82c7aaf36b01cbbbfcc")
+end
+
+getfile(r, name) = r.files[findfirst(x -> x.name == name, r.files)]
+
+function getfile(name)
+  r = ZipFile.Reader(deps("sentiment.zip"))
+  text = read(getfile(r, "trees/$name"), String)
+  close(r)
+  return text
+end
+
+using ..Data: Tree
+
+totree_(n, w) = Tree{Any}((parse(Int, n), w))
+totree_(n, a, b) = Tree{Any}((parse(Int, n), nothing), totree(a), totree(b))
+totree(t::Expr) = totree_(t.args...)
+
+function parsetree(s)
+  s = replace(s, "\\" => "")
+  s = replace(s, "\$" => "\\\$")
+  s = replace(s, r"[^ \n\(\)]+" => s -> "\"$s\"")
+  s = replace(s, " " => ", ")
+  return totree(Meta.parse(s))
+end
+
+function gettrees(name)
+  load()
+  ss = split(getfile("$name.txt"), '\n', keepempty = false)
+  return parsetree.(ss)
+end
+
+"""
+    train()
+
+Return the train split of the Stanford Sentiment Treebank.
+The data is in [treebank](https://en.wikipedia.org/wiki/Treebank) format.
+"""
+train() = gettrees("train")
+
+"""
+    test()
+
+Return the test split of the Stanford Sentiment Treebank.
+The data is in [treebank](https://en.wikipedia.org/wiki/Treebank) format.
+"""
+test() = gettrees("test")
+
+"""
+    dev()
+
+Return the dev split of the Stanford Sentiment Treebank.
+The data is in [treebank](https://en.wikipedia.org/wiki/Treebank) format.
+"""
+dev() = gettrees("dev")
+
+end

src/data/tree.jl

@@ -0,0 +1,42 @@
+using AbstractTrees
+
+struct Tree{T}
+  value::T
+  children::Vector{Tree{T}}
+end
+
+Tree{T}(x::T, xs::Tree{T}...) where T = Tree{T}(x, [xs...])
+Tree{T}(x) where T = Tree(convert(T, x))
+
+Tree(x::T, xs::Tree{T}...) where T = Tree{T}(x, xs...)
+
+AbstractTrees.children(t::Tree) = t.children
+AbstractTrees.printnode(io::IO, t::Tree) = show(io, t.value)
+
+Base.show(io::IO, t::Type{Tree}) = print(io, "Tree")
+Base.show(io::IO, t::Type{Tree{T}}) where T = print(io, "Tree{", T, "}")
+
+function Base.show(io::IO, t::Tree)
+  println(io, typeof(t))
+  print_tree(io, t)
+end
+
+using Juno
+
+@render Juno.Inline t::Tree begin
+  render(t) = Juno.Tree(t.value, render.(t.children))
+  Juno.Tree(typeof(t), [render(t)])
+end
+
+Base.getindex(t::Tree, i::Integer) = t.children[i]
+Base.getindex(t::Tree, i::Integer, is::Integer...) = t[i][is...]
+
+# Utilities
+
+isleaf(t) = isempty(children(t))
+
+leaves(xs::Tree) = map(x -> x.value, Leaves(xs))
+
+Base.map(f, t::Tree, ts::Tree...) =
+  Tree{Any}(f(map(t -> t.value, (t, ts...))...),
+            [map(f, chs...) for chs in zip(map(t -> t.children, (t, ts...))...)]...)

src/deprecations.jl

@@ -0,0 +1,2 @@
+@deprecate param(x) x
+@deprecate data(x) x

src/functor.jl

@@ -0,0 +1,82 @@
+import Adapt: adapt, adapt_storage
+using Zygote: IdSet
+import Functors: @functor, functor, fmap
+
+trainable(m) = functor(m)[1]
+
+"""
+    testmode!(m, mode = true)
+
+Set a layer or model's test mode (see below).
+Using `:auto` mode will treat any gradient computation as training.
+
+_Note_: if you manually set a model into test mode, you need to manually place
+it back into train mode during training phase.
+
+Possible values include:
+- `false` for training
+- `true` for testing
+- `:auto` or `nothing` for Flux to detect the mode automatically
+"""
+testmode!(m, mode = true) = m
+
+"""
+    trainmode!(m, mode = true)
+
+Set a layer of model's train mode (see below).
+Symmetric to [`testmode!`](@ref) (i.e. `trainmode!(m, mode) == testmode!(m, !mode)`).
+
+_Note_: if you manually set a model into train mode, you need to manually place
+it into test mode during testing phase.
+
+Possible values include:
+- `true` for training
+- `false` for testing
+- `:auto` or `nothing` for Flux to detect the mode automatically
+"""
+trainmode!(m, mode = true) = mode isa Bool ? testmode!(m, !mode) : testmode!(m, mode)
+
+params!(p::Params, x::AbstractArray{<:Number}, seen = IdSet()) = push!(p, x)
+
+function params!(p::Params, x, seen = IdSet())
+  x in seen && return
+  push!(seen, x)
+  for child in trainable(x)
+    params!(p, child, seen)
+  end
+end
+
+function params(m...)
+  ps = Params()
+  params!(ps, m)
+  return ps
+end
+
+# Deprecated stuff
+macro treelike(args...)
+  functorm(args...)
+end
+mapleaves(f, x) = fmap(f, x)
+
+function loadparams!(m, xs)
+  for (p, x) in zip(params(m), xs)
+    size(p) == size(x) ||
+      error("Expected param size $(size(p)), got $(size(x))")
+    copyto!(p, x)
+  end
+end
+
+# CPU/GPU movement conveniences
+
+cpu(m) = fmap(x -> adapt(Array, x), m)
+
+gpu(x) = use_cuda[] ? fmap(CuArrays.cu, x) : x
+
+# Precision
+
+adapt_storage(T::Type{<:Real}, xs::AbstractArray{<:Real}) = convert.(T, xs)
+
+paramtype(T::Type{<:Real}, m) = fmap(x -> adapt(T, x), m)
+
+f32(m) = paramtype(Float32, m)
+f64(m) = paramtype(Float64, m)

src/layers/basic.jl

@@ -4,59 +4,120 @@
 Chain multiple layers / functions together, so that they are called in sequence
 on a given input.
 
-    m = Chain(x -> x^2, x -> x+1)
-    m(5) == 26
-
-    m = Chain(Dense(10, 5), Dense(5, 2))
-    x = rand(10)
-    m(x) == m[2](m[1](x))
-
 `Chain` also supports indexing and slicing, e.g. `m[2]` or `m[1:end-1]`.
 `m[1:3](x)` will calculate the output of the first three layers.
+
+# Examples
+```jldoctest
+julia> m = Chain(x -> x^2, x -> x+1);
+
+julia> m(5) == 26
+true
+
+julia> m = Chain(Dense(10, 5), Dense(5, 2));
+
+julia> x = rand(10);
+
+julia> m(x) == m[2](m[1](x))
+true
+```
 """
-type Chain
-  layers::Vector{Any}
-  Chain(xs...) = new([xs...])
+struct Chain{T<:Tuple}
+  layers::T
+  Chain(xs...) = new{typeof(xs)}(xs)
 end
 
-@forward Chain.layers Base.getindex, Base.first, Base.last, Base.endof, Base.push!
-@forward Chain.layers Base.start, Base.next, Base.done
+@forward Chain.layers Base.getindex, Base.length, Base.first, Base.last,
+  Base.iterate, Base.lastindex
 
-children(c::Chain) = c.layers
-mapchildren(f, c::Chain) = Chain(f.(c.layers)...)
+functor(::Type{<:Chain}, c) = c.layers, ls -> Chain(ls...)
 
-(s::Chain)(x) = foldl((x, m) -> m(x), x, s.layers)
+applychain(::Tuple{}, x) = x
+applychain(fs::Tuple, x) = applychain(tail(fs), first(fs)(x))
+
+(c::Chain)(x) = applychain(c.layers, x)
 
 Base.getindex(c::Chain, i::AbstractArray) = Chain(c.layers[i]...)
 
+testmode!(m::Chain, mode = true) = (map(x -> testmode!(x, mode), m.layers); m)
+
 function Base.show(io::IO, c::Chain)
   print(io, "Chain(")
   join(io, c.layers, ", ")
   print(io, ")")
 end
 
+"""
+    outdims(c::Chain, isize)
+
+Calculate the output dimensions given the input dimensions, `isize`.
+
+```julia
+m = Chain(Conv((3, 3), 3 => 16), Conv((3, 3), 16 => 32))
+outdims(m, (10, 10)) == (6, 6)
+```
+"""
+outdims(c::Chain, isize) = foldl(∘, map(l -> (x -> outdims(l, x)), c.layers))(isize)
+
+# This is a temporary and naive implementation
+# it might be replaced in the future for better performance
+# see issue https://github.com/FluxML/Flux.jl/issues/702
+# Johnny Chen -- @johnnychen94
+# only slightly changed to better handle interaction with Zygote @dsweber2
+"""
+    activations(c::Chain, input)
+
+Calculate the forward results of each layers in Chain `c` with `input` as model input.
+"""
+function activations(c::Chain, input)
+    extraChain(c.layers, input)
+end
+
+function extraChain(fs::Tuple, x)
+    res = first(fs)(x)
+    return (res, extraChain(Base.tail(fs), res)...)
+end
+
+extraChain(::Tuple{}, x) = ()
+
+
+
 """
     Dense(in::Integer, out::Integer, σ = identity)
 
-Creates a traditional `Dense` layer with parameters `W` and `b`.
+Create a traditional `Dense` layer with parameters `W` and `b`.
 
     y = σ.(W * x .+ b)
 
 The input `x` must be a vector of length `in`, or a batch of vectors represented
-as an `in × N` matrix. The out `y` will be a vector or batch of length `in`.
+as an `in × N` matrix. The out `y` will be a vector or batch of length `out`.
+
+# Examples
+```jldoctest; setup = :(using Random; Random.seed!(0))
+julia> d = Dense(5, 2)
+Dense(5, 2)
+
+julia> d(rand(5))
+2-element Array{Float32,1}:
+  -0.16210233
+   0.12311903```
 """
-struct Dense{F,S,T}
-  σ::F
+struct Dense{F,S<:AbstractArray,T<:AbstractArray}
   W::S
   b::T
+  σ::F
 end
 
-Dense(in::Integer, out::Integer, σ = identity; init = initn) =
-  Dense(σ, param(init(out, in)), param(init(out)))
+Dense(W, b) = Dense(W, b, identity)
 
-treelike(Dense)
+function Dense(in::Integer, out::Integer, σ = identity;
+               initW = glorot_uniform, initb = zeros)
+  return Dense(initW(out, in), initb(out), σ)
+end
 
-function (a::Dense)(x)
+@functor Dense
+
+function (a::Dense)(x::AbstractArray)
   W, b, σ = a.W, a.b, a.σ
   σ.(W*x .+ b)
 end
@@ -66,3 +127,134 @@ function Base.show(io::IO, l::Dense)
   l.σ == identity || print(io, ", ", l.σ)
   print(io, ")")
 end
+
+# Try to avoid hitting generic matmul in some simple cases
+# Base's matmul is so slow that it's worth the extra conversion to hit BLAS
+(a::Dense{<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  invoke(a, Tuple{AbstractArray}, x)
+
+(a::Dense{<:Any,W})(x::AbstractArray{<:AbstractFloat}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  a(T.(x))
+
+"""
+    outdims(l::Dense, isize)
+
+Calculate the output dimensions given the input dimensions, `isize`.
+
+```julia
+m = Dense(10, 5)
+outdims(m, (5, 2)) == (5,)
+outdims(m, (10,)) == (5,)
+```
+"""
+outdims(l::Dense, isize) = (size(l.W)[1],)
+
+"""
+    Diagonal(in::Integer)
+
+Create an element-wise linear transformation layer with learnable
+vectors `α` and `β`:
+
+    y = α .* x .+ β
+
+The input `x` must be a array where `size(x, 1) == in`.
+"""
+struct Diagonal{T}
+  α::T
+  β::T
+end
+
+Diagonal(in::Integer; initα = ones, initβ = zeros) =
+  Diagonal(initα(in), initβ(in))
+
+@functor Diagonal
+
+function (a::Diagonal)(x)
+  α, β = a.α, a.β
+  α.*x .+ β
+end
+
+function Base.show(io::IO, l::Diagonal)
+  print(io, "Diagonal(", length(l.α), ")")
+end
+
+outdims(l::Diagonal, isize) = (length(l.α),)
+
+"""
+    Maxout(over)
+
+The [Maxout](https://arxiv.org/pdf/1302.4389.pdf) layer has a number of
+internal layers which all receive the same input. It returns the elementwise
+maximum of the internal layers' outputs.
+
+Maxout over linear dense layers satisfies the univeral approximation theorem.
+"""
+struct Maxout{FS<:Tuple}
+    over::FS
+end
+
+"""
+    Maxout(f, n_alts)
+
+Construct a Maxout layer over `n_alts` instances of the layer given by `f`.
+The function takes no arguments and should return some callable layer.
+Conventionally, this is a linear dense layer.
+
+# Examples
+
+This constructs a `Maxout` layer over 4 internal dense linear layers, each
+identical in structure (784 inputs, 128 outputs):
+```julia
+insize = 784
+outsize = 128
+Maxout(()->Dense(insize, outsize), 4)
+```
+"""
+function Maxout(f, n_alts)
+  over = Tuple(f() for _ in 1:n_alts)
+  return Maxout(over)
+end
+
+@functor Maxout
+
+function (mo::Maxout)(input::AbstractArray)
+    mapreduce(f -> f(input), (acc, out) -> max.(acc, out), mo.over)
+end
+
+outdims(l::Maxout, isize) = outdims(first(l.over), isize)
+
+"""
+    SkipConnection(layer, connection)
+
+Create a skip connection which consists of a layer or `Chain` of consecutive
+layers and a shortcut connection linking the block's input to the output
+through a user-supplied 2-argument callable. The first argument to the callable
+will be propagated through the given `layer` while the second is the unchanged,
+"skipped" input.
+
+The simplest "ResNet"-type connection is just `SkipConnection(layer, +)`,
+and requires the output of the layers to be the same shape as the input.
+Here is a more complicated example:
+```julia
+m = Conv((3,3), 4=>7, pad=(1,1))
+x = ones(5,5,4,10);
+size(m(x)) == (5, 5, 7, 10)
+
+sm = SkipConnection(m, (mx, x) -> cat(mx, x, dims=3))
+size(sm(x)) == (5, 5, 11, 10)
+```
+"""
+struct SkipConnection
+  layers
+  connection  #user can pass arbitrary connections here, such as (a,b) -> a + b
+end
+
+@functor SkipConnection
+
+function (skip::SkipConnection)(input)
+  skip.connection(skip.layers(input), input)
+end
+
+function Base.show(io::IO, b::SkipConnection)
+  print(io, "SkipConnection(", b.layers, ", ", b.connection, ")")
+end

src/layers/conv.jl

@@ -0,0 +1,592 @@
+using NNlib: conv, ∇conv_data, depthwiseconv, output_size
+
+# pad dims of x with dims of y until ndims(x) == ndims(y)
+_paddims(x::Tuple, y::Tuple) = (x..., y[(end - (length(y) - length(x) - 1)):end]...)
+
+_convtransoutdims(isize, ksize, ssize, dsize, pad) = (isize .- 1).*ssize .+ 1 .+ (ksize .- 1).*dsize .- (pad[1:2:end] .+ pad[2:2:end])
+
+expand(N, i::Tuple) = i
+expand(N, i::Integer) = ntuple(_ -> i, N)
+
+"""
+    SamePad
+
+Padding for convolutional layers will be calculated so that outputshape == inputshape when stride = 1.
+
+For stride > 1 the output shape depends on the type of convolution layer.
+"""
+struct SamePad end
+
+calc_padding(pad, k::NTuple{N,T}, dilation, stride) where {T,N}= expand(Val(2*N), pad)
+function calc_padding(::SamePad, k::NTuple{N,T}, dilation, stride) where {N,T}
+  #Ref: "A guide to convolution arithmetic for deep learning" https://arxiv.org/pdf/1603.07285
+
+  # Effective kernel size, including dilation
+  k_eff = @. k + (k - 1) * (dilation - 1)
+  # How much total padding needs to be applied?
+  pad_amt = @. k_eff - 1
+  # In case amount of padding is odd we need to apply different amounts to each side.
+  return Tuple(mapfoldl(i -> [ceil(Int, i/2), floor(Int, i/2)], vcat, pad_amt))
+end
+
+"""
+    Conv(filter, in => out, σ = identity; init = glorot_uniform,
+         stride = 1, pad = 0, dilation = 1)
+
+    filter = (2,2)
+    in = 1
+    out = 16
+    Conv((2, 2), 1=>16, relu)
+
+Standard convolutional layer. `filter` should be a tuple like `(2, 2)`.
+`in` and `out` specify the number of input and output channels respectively.
+
+Data should be stored in WHCN order (width, height, # channels, batch size).
+In other words, a 100×100 RGB image would be a `100×100×3×1` array,
+and a batch of 50 would be a `100×100×3×50` array.
+
+Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
+Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
+
+# Examples
+
+Apply a `Conv` layer to a 1-channel input using a 2×2 window filter size, giving us a
+16-channel output. Output is activated with ReLU.
+```julia
+filter = (2,2)
+in = 1
+out = 16
+Conv(filter, in => out, relu)
+```
+"""
+struct Conv{N,M,F,A,V}
+  σ::F
+  weight::A
+  bias::V
+  stride::NTuple{N,Int}
+  pad::NTuple{M,Int}
+  dilation::NTuple{N,Int}
+end
+
+"""
+    Conv(weight::AbstractArray, bias::AbstractArray)
+    Conv(weight::AbstractArray, bias::AbstractArray, activation)
+
+Constructs the convolutional layer with user defined weight and bias arrays.
+
+Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+
+There is also a keyword-only constuctor available for all convoultional
+layers.
+
+```julia
+weight = rand(Float32, 3, 3, 5)
+bias = zeros(Float32, 5)
+Conv(weight = weight,
+    bias = bias,
+    σ = sigmoid)
+```
+"""
+function Conv(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
+              stride = 1, pad = 0, dilation = 1) where {T,N}
+  stride = expand(Val(N-2), stride)
+  dilation = expand(Val(N-2), dilation)
+  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
+  return Conv(σ, w, b, stride, pad, dilation)
+end
+
+function Conv(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
+              activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
+  Conv(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
+end
+
+"""
+    convfilter(filter::Tuple, in=>out)
+
+Constructs a standard convolutional weight matrix with given `filter` and
+channels from `in` to `out`.
+
+Accepts the keyword `init` (default: `glorot_uniform`) to control the sampling
+distribution.
+
+See also: [`depthwiseconvfilter`](@ref)
+"""
+convfilter(filter::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer};
+          init = glorot_uniform) where N = init(filter..., ch...)
+
+function Conv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
+            init = glorot_uniform,  stride = 1, pad = 0, dilation = 1,
+            weight = convfilter(k, ch, init = init), bias = zeros(ch[2])) where N
+
+  Conv(weight, bias, σ,
+      stride = stride, pad = pad, dilation = dilation)
+end
+
+@functor Conv
+
+function (c::Conv)(x::AbstractArray)
+  # TODO: breaks gpu broadcast :(
+  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
+  σ, b = c.σ, reshape(c.bias, ntuple(_->1, length(c.stride))..., :, 1)
+  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
+  σ.(conv(x, c.weight, cdims) .+ b)
+end
+
+function Base.show(io::IO, l::Conv)
+  print(io, "Conv(", size(l.weight)[1:ndims(l.weight)-2])
+  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
+  l.σ == identity || print(io, ", ", l.σ)
+  print(io, ")")
+end
+
+(a::Conv{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  invoke(a, Tuple{AbstractArray}, x)
+
+(a::Conv{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  a(T.(x))
+
+"""
+    outdims(l::Conv, isize::Tuple)
+
+Calculate the output dimensions given the input dimensions `isize`.
+Batch size and channel size are ignored as per [NNlib.jl](https://github.com/FluxML/NNlib.jl).
+
+```julia
+m = Conv((3, 3), 3 => 16)
+outdims(m, (10, 10)) == (8, 8)
+outdims(m, (10, 10, 1, 3)) == (8, 8)
+```
+"""
+outdims(l::Conv, isize) =
+  output_size(DenseConvDims(_paddims(isize, size(l.weight)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))
+
+"""
+    ConvTranspose(filter, in=>out)
+    ConvTranspose(filter, in=>out, activation)
+    ConvTranspose(filter, in => out, σ = identity; init = glorot_uniform,
+                  stride = 1, pad = 0, dilation = 1)
+
+Standard convolutional transpose layer. `filter` should be a tuple like `(2, 2)`.
+`in` and `out` specify the number of input and output channels respectively.
+
+Data should be stored in WHCN order (width, height, # channels, batch size).
+In other words, a 100×100 RGB image would be a `100×100×3×1` array,
+and a batch of 50 would be a `100×100×3×50` array.
+
+Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
+Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+Use `pad=SamePad()` to apply padding so that outputsize == stride * inputsize - stride + 1.
+"""
+struct ConvTranspose{N,M,F,A,V}
+  σ::F
+  weight::A
+  bias::V
+  stride::NTuple{N,Int}
+  pad::NTuple{M,Int}
+  dilation::NTuple{N,Int}
+end
+
+"""
+    ConvTranspose(weight::AbstractArray, bias::AbstractArray)
+    ConvTranspose(weight::AbstractArray, bias::AbstractArray, activation)
+
+Constructs the convolutional transpose layer with user defined weight and bias arrays.
+forward pass.
+
+Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+
+For keyword-only constuctor, see also [`Conv`](@ref)
+"""
+function ConvTranspose(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
+                      stride = 1, pad = 0, dilation = 1) where {T,N}
+  stride = expand(Val(N-2), stride)
+  dilation = expand(Val(N-2), dilation)
+  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
+  return ConvTranspose(σ, w, b, stride, pad, dilation)
+end
+
+function ConvTranspose(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
+                        activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
+  ConvTranspose(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
+end
+
+function ConvTranspose(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
+                      init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
+                      weight = convfilter(k, reverse(ch), init = init), bias = zeros(ch[2])) where N
+
+  ConvTranspose(weight, bias, σ,
+              stride = stride, pad = pad, dilation = dilation)
+end
+
+@functor ConvTranspose
+
+function conv_transpose_dims(c::ConvTranspose, x::AbstractArray)
+    # Calculate size of "input", from ∇conv_data()'s perspective...
+    combined_pad = (c.pad[1:2:end] .+ c.pad[2:2:end])
+    I = (size(x)[1:end-2] .- 1).*c.stride .+ 1 .+ (size(c.weight)[1:end-2] .- 1).*c.dilation .- combined_pad
+    C_in = size(c.weight)[end-1]
+    batch_size = size(x)[end]
+    # Create DenseConvDims() that looks like the corresponding conv()
+    return DenseConvDims((I..., C_in, batch_size), size(c.weight);
+                        stride=c.stride,
+                        padding=c.pad,
+                        dilation=c.dilation,
+    )
+end
+
+# TODO: Find proper fix for https://github.com/FluxML/Flux.jl/issues/900
+@nograd conv_transpose_dims
+
+function (c::ConvTranspose)(x::AbstractArray)
+  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
+  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
+  cdims = conv_transpose_dims(c, x)
+  σ.(∇conv_data(x, c.weight, cdims) .+ b)
+end
+
+function Base.show(io::IO, l::ConvTranspose)
+  print(io, "ConvTranspose(", size(l.weight)[1:ndims(l.weight)-2])
+  print(io, ", ", size(l.weight, ndims(l.weight)), "=>", size(l.weight, ndims(l.weight)-1))
+  l.σ == identity || print(io, ", ", l.σ)
+  print(io, ")")
+end
+
+(a::ConvTranspose{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  invoke(a, Tuple{AbstractArray}, x)
+
+(a::ConvTranspose{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  a(T.(x))
+
+outdims(l::ConvTranspose{N}, isize) where N = _convtransoutdims(isize[1:2], size(l.weight)[1:N], l.stride, l.dilation, l.pad)
+
+"""
+    DepthwiseConv(filter::Tuple, in=>out)
+    DepthwiseConv(filter::Tuple, in=>out, activation)
+    DepthwiseConv(filter, in => out, σ = identity; init = glorot_uniform,
+                  stride = 1, pad = 0, dilation = 1)
+
+Depthwise convolutional layer. `filter` should be a tuple like `(2, 2)`.
+`in` and `out` specify the number of input and output channels respectively.
+Note that `out` must be an integer multiple of `in`.
+
+Data should be stored in WHCN order (width, height, # channels, batch size).
+In other words, a 100×100 RGB image would be a `100×100×3×1` array,
+and a batch of 50 would be a `100×100×3×50` array.
+
+Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
+Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
+"""
+struct DepthwiseConv{N,M,F,A,V}
+  σ::F
+  weight::A
+  bias::V
+  stride::NTuple{N,Int}
+  pad::NTuple{M,Int}
+  dilation::NTuple{N,Int}
+end
+
+"""
+    DepthwiseConv(weight::AbstractArray, bias::AbstractArray)
+    DepthwiseConv(weight::AbstractArray, bias::AbstractArray, activation)
+
+Constructs the `DepthwiseConv` layer with user defined weight and bias arrays.
+forward pass.
+
+Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+
+For keyword-only constuctor, see also [`Conv`](@ref)
+"""
+function DepthwiseConv(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
+                      stride = 1, pad = 0, dilation = 1) where {T,N}
+  stride = expand(Val(N-2), stride)
+  dilation = expand(Val(N-2), dilation)
+  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
+  return DepthwiseConv(σ, w, b, stride, pad, dilation)
+end
+
+function DepthwiseConv(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
+                      activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
+  DepthwiseConv(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
+end
+
+"""
+    depthwiseconvfilter(filter::Tuple, in=>out)
+
+Constructs a depthwise convolutional weight array defined by `filter` and channels
+from `in` to `out`.
+
+Accepts the keyword `init` (default: `glorot_uniform`) to control the sampling
+distribution.
+
+See also: [`convfilter`](@ref)
+"""
+depthwiseconvfilter(filter::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer};
+                    init = glorot_uniform) where N = init(filter..., div(ch[2], ch[1]), ch[1])
+
+function DepthwiseConv(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
+                      init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
+                      weight = depthwiseconvfilter(k, ch, init = init), bias = zeros(ch[2])) where N
+  @assert ch[2] % ch[1] == 0 "Output channels must be integer multiple of input channels"
+
+  return DepthwiseConv(
+    weight,
+    bias,
+    σ;
+    stride = stride,
+    pad = pad,
+    dilation = dilation
+  )
+end
+
+@functor DepthwiseConv
+
+function (c::DepthwiseConv)(x)
+  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
+  cdims = DepthwiseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
+  σ.(depthwiseconv(x, c.weight, cdims) .+ b)
+end
+
+function Base.show(io::IO, l::DepthwiseConv)
+  print(io, "DepthwiseConv(", size(l.weight)[1:end-2])
+  print(io, ", ", size(l.weight)[end], "=>", prod(size(l.weight)[end-1:end]))
+  l.σ == identity || print(io, ", ", l.σ)
+  print(io, ")")
+end
+
+(a::DepthwiseConv{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  invoke(a, Tuple{AbstractArray}, x)
+
+(a::DepthwiseConv{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  a(T.(x))
+
+outdims(l::DepthwiseConv, isize) =
+  output_size(DepthwiseConvDims(_paddims(isize, (1, 1, size(l.weight)[end], 1)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))
+
+"""
+    CrossCor(filter, in=>out)
+    CrossCor(filter, in=>out, activation)
+    CrossCor(filter, in => out, σ = identity; init = glorot_uniform,
+             stride = 1, pad = 0, dilation = 1)
+
+Standard cross convolutional layer. `filter` should be a tuple like `(2, 2)`.
+`in` and `out` specify the number of input and output channels respectively.
+
+Data should be stored in WHCN order (width, height, # channels, batch size).
+In other words, a 100×100 RGB image would be a `100×100×3×1` array,
+and a batch of 50 would be a `100×100×3×50` array.
+
+Accepts keyword arguments `weight` and `bias` to set the corresponding fields.
+Setting `bias` to `Flux.Zeros()` will switch bias off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
+
+# Examples
+
+Apply a `CrossCor` layer to a 1-channel input using a 2×2 window filter size, giving us a
+16-channel output. Output is activated with ReLU.
+```julia
+filter = (2,2)
+in = 1
+out = 16
+CrossCor((2, 2), 1=>16, relu)
+```
+"""
+struct CrossCor{N,M,F,A,V}
+  σ::F
+  weight::A
+  bias::V
+  stride::NTuple{N,Int}
+  pad::NTuple{M,Int}
+  dilation::NTuple{N,Int}
+end
+
+"""
+    CrossCor(weight::AbstractArray, bias::AbstractArray)
+    CrossCor(weight::AbstractArray, bias::AbstractArray, activation)
+
+Constructs the standard cross convolutional layer with user defined weight and bias
+arrays.
+
+Setting `bias` to `Flux.Zeros()` would switch `bias` off for the layer.
+
+Takes the keyword arguments `pad`, `stride` and `dilation`.
+
+For keyword-only constuctor, see also [`Conv`](@ref)
+"""
+function CrossCor(w::AbstractArray{T,N}, b::Union{Zeros, AbstractVector{T}}, σ = identity;
+                  stride = 1, pad = 0, dilation = 1) where {T,N}
+  stride = expand(Val(N-2), stride)
+  dilation = expand(Val(N-2), dilation)
+  pad = calc_padding(pad, size(w)[1:N-2], dilation, stride)
+  return CrossCor(σ, w, b, stride, pad, dilation)
+end
+
+function CrossCor(;weight::AbstractArray{T,N}, bias::Union{Zeros, AbstractVector{T}},
+                      activation = identity, stride = 1, pad = 0, dilation = 1) where {T,N}
+  CrossCor(weight, bias, activation, stride = stride, pad = pad, dilation = dilation)
+end
+
+function CrossCor(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
+                  init = glorot_uniform, stride = 1, pad = 0, dilation = 1,
+                  weight = convfilter(k, ch, init = init), bias = zeros(ch[2])) where N
+
+  CrossCor(weight, bias, σ,
+       stride = stride, pad = pad, dilation = dilation)
+end
+
+@functor CrossCor
+
+function crosscor(x, w, ddims::DenseConvDims)
+  ddims = DenseConvDims(ddims, F=true)
+  return conv(x, w, ddims)
+end
+
+function (c::CrossCor)(x::AbstractArray)
+  # TODO: breaks gpu broadcast :(
+  # ndims(x) == ndims(c.weight)-1 && return squeezebatch(c(reshape(x, size(x)..., 1)))
+  σ, b = c.σ, reshape(c.bias, map(_->1, c.stride)..., :, 1)
+  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
+  σ.(crosscor(x, c.weight, cdims) .+ b)
+end
+
+function Base.show(io::IO, l::CrossCor)
+  print(io, "CrossCor(", size(l.weight)[1:ndims(l.weight)-2])
+  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
+  l.σ == identity || print(io, ", ", l.σ)
+  print(io, ")")
+end
+
+(a::CrossCor{<:Any,<:Any,W})(x::AbstractArray{T}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  invoke(a, Tuple{AbstractArray}, x)
+
+(a::CrossCor{<:Any,<:Any,W})(x::AbstractArray{<:Real}) where {T <: Union{Float32,Float64}, W <: AbstractArray{T}} =
+  a(T.(x))
+
+outdims(l::CrossCor, isize) =
+  output_size(DenseConvDims(_paddims(isize, size(l.weight)), size(l.weight); stride = l.stride, padding = l.pad, dilation = l.dilation))
+
+"""
+    GlobalMaxPool()
+
+Global max pooling layer.
+
+Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output,
+by performing max pooling on the complete (w,h)-shaped feature maps.
+"""
+struct GlobalMaxPool end
+
+function (g::GlobalMaxPool)(x)
+  # Input size
+  x_size = size(x)
+  # Kernel size
+  k = x_size[1:end-2]
+  # Pooling dimensions
+  pdims = PoolDims(x, k)
+
+  return maxpool(x, pdims)
+end
+
+function Base.show(io::IO, g::GlobalMaxPool)
+  print(io, "GlobalMaxPool()")
+end
+
+"""
+    GlobalMeanPool()
+
+Global mean pooling layer.
+
+Transforms (w,h,c,b)-shaped input into (1,1,c,b)-shaped output,
+by performing mean pooling on the complete (w,h)-shaped feature maps.
+"""
+struct GlobalMeanPool end
+
+function (g::GlobalMeanPool)(x)
+  # Input size
+  x_size = size(x)
+  # Kernel size
+  k = x_size[1:end-2]
+  # Pooling dimensions
+  pdims = PoolDims(x, k)
+
+  return meanpool(x, pdims)
+end
+
+function Base.show(io::IO, g::GlobalMeanPool)
+  print(io, "GlobalMeanPool()")
+end
+
+"""
+    MaxPool(k; pad = 0, stride = k)
+
+Max pooling layer. `k` is the size of the window for each dimension of the input.
+
+Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
+=======
+"""
+struct MaxPool{N,M}
+  k::NTuple{N,Int}
+  pad::NTuple{M,Int}
+  stride::NTuple{N,Int}
+end
+
+function MaxPool(k::NTuple{N,Integer}; pad = 0, stride = k) where N
+  stride = expand(Val(N), stride)
+  pad = calc_padding(pad, k, 1, stride)
+  return MaxPool(k, pad, stride)
+end
+
+function (m::MaxPool)(x)
+    pdims = PoolDims(x, m.k; padding=m.pad, stride=m.stride)
+    return maxpool(x, pdims)
+end
+
+function Base.show(io::IO, m::MaxPool)
+  print(io, "MaxPool(", m.k, ", pad = ", m.pad, ", stride = ", m.stride, ")")
+end
+
+outdims(l::MaxPool{N}, isize) where N = output_size(PoolDims(_paddims(isize, (l.k..., 1, 1)), l.k; stride = l.stride, padding = l.pad))
+
+"""
+    MeanPool(k; pad = 0, stride = k)
+
+Mean pooling layer. `k` is the size of the window for each dimension of the input.
+
+Use `pad=SamePad()` to apply padding so that outputsize == inputsize / stride.
+"""
+struct MeanPool{N,M}
+    k::NTuple{N,Int}
+    pad::NTuple{M,Int}
+    stride::NTuple{N,Int}
+end
+
+function MeanPool(k::NTuple{N,Integer}; pad = 0, stride = k) where N
+  stride = expand(Val(N), stride)
+  pad = calc_padding(pad, k, 1, stride)
+  return MeanPool(k, pad, stride)
+end
+
+function (m::MeanPool)(x)
+    pdims = PoolDims(x, m.k; padding=m.pad, stride=m.stride)
+    return meanpool(x, pdims)
+end
+
+function Base.show(io::IO, m::MeanPool)
+  print(io, "MeanPool(", m.k, ", pad = ", m.pad, ", stride = ", m.stride, ")")
+end
+
+outdims(l::MeanPool{N}, isize) where N = output_size(PoolDims(_paddims(isize, (l.k..., 1, 1)), l.k; stride = l.stride, padding = l.pad))

src/layers/normalise.jl

@@ -0,0 +1,416 @@
+istraining() = false
+
+@adjoint istraining() = true, _ -> nothing
+
+_isactive(m) = isnothing(m.active) ? istraining() : m.active
+
+_dropout_shape(s, ::Colon) = size(s)
+_dropout_shape(s, dims) = tuple((i ∉ dims ? 1 : si for (i, si) ∈ enumerate(size(s)))...)
+
+_dropout_kernel(y::T, p, q) where {T} = y > p ? T(1 / q) : T(0)
+
+"""
+    dropout(x, p; dims = :)
+
+The dropout function. For each input, either sets that input to `0` (with probability
+`p`) or scales it by `1 / (1 - p)`. `dims` specifies the unbroadcasted dimensions,
+e.g. `dims=1` applies dropout along columns and `dims=2` along rows.
+This is used as a regularisation, i.e. it reduces overfitting during training.
+
+See also the [`Dropout`](@ref) layer.
+"""
+dropout(x, p; dims = :) = x
+
+@adjoint function dropout(x, p; dims = :)
+  y = rand!(similar(x, _dropout_shape(x, dims)))
+  y .= _dropout_kernel.(y, p, 1 - p)
+  return x .* y, Δ -> (Δ .* y, nothing)
+end
+
+"""
+    Dropout(p, dims = :)
+
+Dropout layer. In the forward pass, apply the [`Flux.dropout`](@ref) function on the input.
+
+Does nothing to the input once [`Flux.testmode!`](@ref) is `true`.
+"""
+mutable struct Dropout{F,D}
+  p::F
+  dims::D
+  active::Union{Bool, Nothing}
+end
+
+# TODO: deprecate in v0.11
+Dropout(p, dims) = Dropout(p, dims, nothing)
+
+function Dropout(p; dims = :)
+  @assert 0 ≤ p ≤ 1
+  Dropout{typeof(p),typeof(dims)}(p, dims, nothing)
+end
+
+function (a::Dropout)(x)
+  _isactive(a) || return x
+  return dropout(x, a.p; dims = a.dims)
+end
+
+testmode!(m::Dropout, mode = true) =
+  (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+
+function Base.show(io::IO, d::Dropout)
+  print(io, "Dropout(", d.p)
+  d.dims != (:) && print(io, ", dims = $(repr(d.dims))")
+  print(io, ")")
+end
+
+"""
+    AlphaDropout(p)
+
+A dropout layer. Used in
+[Self-Normalizing Neural Networks](https://papers.nips.cc/paper/6698-self-normalizing-neural-networks.pdf).
+The AlphaDropout layer ensures that mean and variance of activations
+remain the same as before.
+
+Does nothing to the input once [`testmode!`](@ref) is true.
+"""
+mutable struct AlphaDropout{F}
+  p::F
+  active::Union{Bool, Nothing}
+  function AlphaDropout(p, active = nothing)
+    @assert 0 ≤ p ≤ 1
+    new{typeof(p)}(p, active)
+  end
+end
+
+function (a::AlphaDropout)(x)
+  _isactive(a) || return x
+  λ = eltype(x)(1.0507009873554804934193349852946)
+  α = eltype(x)(1.6732632423543772848170429916717)
+  α1 = eltype(x)(-λ*α)
+  noise = randn(eltype(x), size(x))
+  x = @. x*(noise > (1 - a.p)) + α1 * (noise < (1 - a.p))
+  A = (a.p + a.p * (1 - a.p) * α1 ^ 2)^0.5
+  B = -A * α1 * (1 - a.p)
+  x = @. A * x + B
+  return x
+end
+
+testmode!(m::AlphaDropout, mode = true) =
+  (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+
+"""
+    LayerNorm(h::Integer)
+
+A [normalisation layer](https://arxiv.org/pdf/1607.06450.pdf) designed to be
+used with recurrent hidden states of size `h`. Normalises the mean and standard
+deviation of each input before applying a per-neuron gain/bias.
+"""
+struct LayerNorm{T}
+  diag::Diagonal{T}
+end
+
+LayerNorm(h::Integer) =
+  LayerNorm(Diagonal(h))
+
+@functor LayerNorm
+
+(a::LayerNorm)(x) = a.diag(normalise(x))
+
+function Base.show(io::IO, l::LayerNorm)
+  print(io, "LayerNorm(", length(l.diag.α), ")")
+end
+
+"""
+    BatchNorm(channels::Integer, σ = identity;
+              initβ = zeros, initγ = ones,
+              ϵ = 1e-8, momentum = .1)
+
+[Batch Normalization](https://arxiv.org/pdf/1502.03167.pdf) layer.
+`channels` should be the size of the channel dimension in your data (see below).
+
+Given an array with `N` dimensions, call the `N-1`th the channel dimension. (For
+a batch of feature vectors this is just the data dimension, for `WHCN` images
+it's the usual channel dimension.)
+
+`BatchNorm` computes the mean and variance for each each `W×H×1×N` slice and
+shifts them to have a new mean and variance (corresponding to the learnable,
+per-channel `bias` and `scale` parameters).
+
+Use [`testmode!`](@ref) during inference.
+
+# Examples
+```julia
+m = Chain(
+  Dense(28^2, 64),
+  BatchNorm(64, relu),
+  Dense(64, 10),
+  BatchNorm(10),
+  softmax)
+```
+"""
+mutable struct BatchNorm{F,V,W,N}
+  λ::F  # activation function
+  β::V  # bias
+  γ::V  # scale
+  μ::W  # moving mean
+  σ²::W  # moving std
+  ϵ::N
+  momentum::N
+  active::Union{Bool, Nothing}
+end
+
+# TODO: deprecate in v0.11
+BatchNorm(λ, β, γ, μ, σ², ϵ, momentum) = BatchNorm(λ, β, γ, μ, σ², ϵ, momentum, nothing)
+
+BatchNorm(chs::Integer, λ = identity;
+          initβ = (i) -> zeros(Float32, i), initγ = (i) -> ones(Float32, i), ϵ = 1f-5, momentum = 0.1f0) =
+  BatchNorm(λ, initβ(chs), initγ(chs),
+            zeros(chs), ones(chs), ϵ, momentum, nothing)
+
+trainable(bn::BatchNorm) = (bn.β, bn.γ)
+
+function (BN::BatchNorm)(x)
+  size(x, ndims(x)-1) == length(BN.β) ||
+    error("BatchNorm expected $(length(BN.β)) channels, got $(size(x, ndims(x)-1))")
+  dims = length(size(x))
+  channels = size(x, dims-1)
+  affine_shape = ntuple(i->i == ndims(x) - 1 ? size(x, i) : 1, ndims(x))
+  m = div(prod(size(x)), channels)
+  γ = reshape(BN.γ, affine_shape...)
+  β = reshape(BN.β, affine_shape...)
+  if !_isactive(BN)
+    μ = reshape(BN.μ, affine_shape...)
+    σ² = reshape(BN.σ², affine_shape...)
+    ϵ = BN.ϵ
+  else
+    T = eltype(x)
+    axes = [1:dims-2; dims] # axes to reduce along (all but channels axis)
+    μ = mean(x, dims = axes)
+    σ² = sum((x .- μ) .^ 2, dims = axes) ./ m
+    ϵ = convert(T, BN.ϵ)
+    # update moving mean/std
+    mtm = BN.momentum
+    S = eltype(BN.μ)
+    BN.μ  = (1 - mtm) .* BN.μ .+ mtm .* S.(reshape(μ, :))
+    BN.σ² = (1 - mtm) .* BN.σ² .+ (mtm * m / (m - 1)) .* S.(reshape(σ², :))
+  end
+
+  let λ = BN.λ
+    x̂ = (x .- μ) ./ sqrt.(σ² .+ ϵ)
+    λ.(γ .* x̂ .+ β)
+  end
+end
+
+@functor BatchNorm
+
+testmode!(m::BatchNorm, mode = true) =
+  (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+
+function Base.show(io::IO, l::BatchNorm)
+  print(io, "BatchNorm($(join(size(l.β), ", "))")
+  (l.λ == identity) || print(io, ", λ = $(l.λ)")
+  print(io, ")")
+end
+
+expand_inst = (x, as) -> reshape(repeat(x, outer=[1, as[length(as)]]), as...)
+
+mutable struct InstanceNorm{F,V,W,N}
+  λ::F  # activation function
+  β::V  # bias
+  γ::V  # scale
+  μ::W  # moving mean
+  σ²::W  # moving std
+  ϵ::N
+  momentum::N
+  active::Union{Bool, Nothing}
+end
+
+# TODO: deprecate in v0.11
+"""
+    InstanceNorm(channels::Integer, σ = identity;
+                 initβ = zeros, initγ = ones,
+                 ϵ = 1e-8, momentum = .1)
+
+[Instance Normalization](https://arxiv.org/abs/1607.08022) layer.
+`channels` should be the size of the channel dimension in your data (see below).
+
+Given an array with `N` dimensions, call the `N-1`th the channel dimension. (For
+a batch of feature vectors this is just the data dimension, for `WHCN` images
+it's the usual channel dimension.)
+
+`InstanceNorm` computes the mean and variance for each each `W×H×1×1` slice and
+shifts them to have a new mean and variance (corresponding to the learnable,
+per-channel `bias` and `scale` parameters).
+
+Use [`testmode!`](@ref) during inference.
+
+# Examples
+```julia
+m = Chain(
+  Dense(28^2, 64),
+  InstanceNorm(64, relu),
+  Dense(64, 10),
+  InstanceNorm(10),
+  softmax)
+```
+"""
+InstanceNorm(λ, β, γ, μ, σ², ϵ, momentum) = InstanceNorm(λ, β, γ, μ, σ², ϵ, momentum, nothing)
+
+InstanceNorm(chs::Integer, λ = identity;
+          initβ = (i) -> zeros(Float32, i), initγ = (i) -> ones(Float32, i), ϵ = 1f-5, momentum = 0.1f0) =
+  InstanceNorm(λ, initβ(chs), initγ(chs),
+            zeros(chs), ones(chs), ϵ, momentum, nothing)
+
+trainable(in::InstanceNorm) = (in.β, in.γ)
+
+function (in::InstanceNorm)(x)
+  size(x, ndims(x)-1) == length(in.β) ||
+    error("InstanceNorm expected $(length(in.β)) channels, got $(size(x, ndims(x)-1))")
+  ndims(x) > 2 ||
+    error("InstanceNorm requires at least 3 dimensions. With 2 dimensions an array of zeros would be returned")
+  # these are repeated later on depending on the batch size
+  dims = length(size(x))
+  c = size(x, dims-1)
+  bs = size(x, dims)
+  affine_shape = ntuple(i->i == ndims(x) - 1 || i == ndims(x) ? size(x, i) : 1, ndims(x))
+  m = div(prod(size(x)), c*bs)
+  γ, β = expand_inst(in.γ, affine_shape), expand_inst(in.β, affine_shape)
+
+  if !_isactive(in)
+    μ = expand_inst(in.μ, affine_shape)
+    σ² = expand_inst(in.σ², affine_shape)
+    ϵ = in.ϵ
+  else
+    T = eltype(x)
+
+    ϵ = convert(T, in.ϵ)
+    axes = 1:dims-2 # axes to reduce along (all but channels and batch size axes)
+    μ = mean(x, dims = axes)
+    σ² = mean((x .- μ) .^ 2, dims = axes)
+    S = eltype(in.μ)
+    # update moving mean/std
+    mtm = in.momentum
+    in.μ = dropdims(mean(repeat((1 - mtm) .* in.μ, outer=[1, bs]) .+ mtm .* S.(reshape(μ, (c, bs))), dims = 2), dims=2)
+    in.σ² = dropdims(mean((repeat((1 - mtm) .* in.σ², outer=[1, bs]) .+ (mtm * m / (m - 1)) .* S.(reshape(σ², (c, bs)))), dims = 2), dims=2)
+  end
+
+  let λ = in.λ
+    x̂ = (x .- μ) ./ sqrt.(σ² .+ ϵ)
+    λ.(γ .* x̂ .+ β)
+  end
+end
+
+@functor InstanceNorm
+
+testmode!(m::InstanceNorm, mode = true) =
+  (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+
+function Base.show(io::IO, l::InstanceNorm)
+  print(io, "InstanceNorm($(join(size(l.β), ", "))")
+  (l.λ == identity) || print(io, ", λ = $(l.λ)")
+  print(io, ")")
+end
+
+"""
+    GroupNorm(chs::Integer, G::Integer, λ = identity;
+              initβ = (i) -> zeros(Float32, i), initγ = (i) -> ones(Float32, i),
+              ϵ = 1f-5, momentum = 0.1f0)
+
+[Group Normalization](https://arxiv.org/pdf/1803.08494.pdf) layer.
+This layer can outperform Batch Normalization and Instance Normalization.
+
+`chs` is the number of channels, the channel dimension of your input.
+For an array of N dimensions, the `N-1`th index is the channel dimension.
+
+`G` is the number of groups along which the statistics are computed.
+The number of channels must be an integer multiple of the number of groups.
+
+Use [`testmode!`](@ref) during inference.
+
+# Examples
+```julia
+m = Chain(Conv((3,3), 1=>32, leakyrelu;pad = 1),
+          GroupNorm(32,16))
+          # 32 channels, 16 groups (G = 16), thus 2 channels per group used
+```
+"""
+mutable struct GroupNorm{F,V,W,N,T}
+  G::T # number of groups
+  λ::F  # activation function
+  β::V  # bias
+  γ::V  # scale
+  μ::W  # moving mean
+  σ²::W  # moving std
+  ϵ::N
+  momentum::N
+  active::Union{Bool, Nothing}
+end
+
+# TODO: deprecate in v0.11
+GroupNorm(G, λ, β, γ, μ, σ², ϵ, momentum) = GroupNorm(G, λ, β, γ, μ, σ², ϵ, momentum, nothing)
+
+GroupNorm(chs::Integer, G::Integer, λ = identity;
+          initβ = (i) -> zeros(Float32, i), initγ = (i) -> ones(Float32, i), ϵ = 1f-5, momentum = 0.1f0) =
+  GroupNorm(G, λ, initβ(chs), initγ(chs),
+            zeros(G,1), ones(G,1), ϵ, momentum, nothing)
+
+trainable(gn::GroupNorm) = (gn.β, gn.γ)
+
+function(gn::GroupNorm)(x)
+  size(x,ndims(x)-1) == length(gn.β) || error("Group Norm expected $(length(gn.β)) channels, but got $(size(x,ndims(x)-1)) channels")
+  ndims(x) > 2 || error("Need to pass at least 3 channels for Group Norm to work")
+  (size(x,ndims(x) -1))%gn.G == 0 || error("The number of groups ($(gn.G)) must divide the number of channels ($(size(x,ndims(x) -1)))")
+
+  dims = length(size(x))
+  groups = gn.G
+  channels = size(x, dims-1)
+  batches = size(x,dims)
+  channels_per_group = div(channels,groups)
+  affine_shape = ntuple(i->i == ndims(x) - 1 ? size(x, i) : 1, ndims(x))
+
+  # Output reshaped to (W,H...,C/G,G,N)
+  μ_affine_shape = ntuple(i->i == ndims(x) ? groups : 1, ndims(x) + 1)
+
+  m = prod(size(x)[1:end-2]) * channels_per_group
+  γ = reshape(gn.γ, affine_shape...)
+  β = reshape(gn.β, affine_shape...)
+
+  y = reshape(x,((size(x))[1:end-2]...,channels_per_group,groups,batches))
+  if !_isactive(gn)
+    og_shape = size(x)
+    μ = reshape(gn.μ, μ_affine_shape...) # Shape : (1,1,...C/G,G,1)
+    σ² = reshape(gn.σ², μ_affine_shape...) # Shape : (1,1,...C/G,G,1)
+    ϵ = gn.ϵ
+  else
+    T = eltype(x)
+    og_shape = size(x)
+    axes = [(1:ndims(y)-2)...] # axes to reduce along (all but channels axis)
+    μ = mean(y, dims = axes)
+    σ² = mean((y .- μ) .^ 2, dims = axes)
+
+    ϵ = convert(T, gn.ϵ)
+    # update moving mean/std
+    mtm = gn.momentum
+    S = eltype(gn.μ)
+    gn.μ = mean((1 - mtm) .* gn.μ .+ mtm .* S.(reshape(μ, (groups,batches))),dims=2)
+    gn.σ² = mean((1 - mtm) .* gn.σ² .+ (mtm * m / (m - 1)) .* S.(reshape(σ², (groups,batches))),dims=2)
+  end
+
+  let λ = gn.λ
+    x̂ = (y .- μ) ./ sqrt.(σ² .+ ϵ)
+
+    # Reshape x̂
+    x̂ = reshape(x̂,og_shape)
+    λ.(γ .* x̂ .+ β)
+  end
+end
+
+@functor GroupNorm
+
+testmode!(m::GroupNorm, mode = true) =
+  (m.active = (isnothing(mode) || mode == :auto) ? nothing : !mode; m)
+
+function Base.show(io::IO, l::GroupNorm)
+  print(io, "GroupNorm($(join(size(l.β), ", "))")
+  (l.λ == identity) || print(io, ", λ = $(l.λ)")
+  print(io, ")")
+end

src/layers/recurrent.jl

@@ -1,14 +1,36 @@
-# TODO: broadcasting cat
-combine(x, h) = vcat(x, h .* trues(1, size(x, 2)))
+gate(h, n) = (1:h) .+ h*(n-1)
+gate(x::AbstractVector, h, n) = @view x[gate(h,n)]
+gate(x::AbstractMatrix, h, n) = x[gate(h,n),:]
 
 # Stateful recurrence
 
+"""
+    Recur(cell)
+
+`Recur` takes a recurrent cell and makes it stateful, managing the hidden state
+in the background. `cell` should be a model of the form:
+
+    h, y = cell(h, x...)
+
+For example, here's a recurrent network that keeps a running total of its inputs:
+
+```julia
+accum(h, x) = (h + x, x)
+rnn = Flux.Recur(accum, 0)
+rnn(2)      # 2
+rnn(3)      # 3
+rnn.state   # 5
+rnn.(1:10)  # apply to a sequence
+rnn.state   # 60
+```
+"""
 mutable struct Recur{T}
   cell::T
+  init
   state
 end
 
-Recur(m) = Recur(m, hidden(m))
+Recur(m, h = hidden(m)) = Recur(m, h, h)
 
 function (m::Recur)(xs...)
   h, y = m.cell(m.state, xs...)
@@ -16,79 +38,149 @@ function (m::Recur)(xs...)
   return y
 end
 
-treelike(Recur)
+@functor Recur cell, init
 
 Base.show(io::IO, m::Recur) = print(io, "Recur(", m.cell, ")")
 
-_truncate(x::AbstractArray) = x
-_truncate(x::TrackedArray) = x.data
-_truncate(x::Tuple) = _truncate.(x)
+"""
+    reset!(rnn)
 
-truncate!(m) = foreach(truncate!, children(m))
-truncate!(m::Recur) = (m.state = _truncate(m.state))
+Reset the hidden state of a recurrent layer back to its original value.
+
+Assuming you have a `Recur` layer `rnn`, this is roughly equivalent to:
+```julia
+rnn.state = hidden(rnn.cell)
+```
+"""
+reset!(m::Recur) = (m.state = m.init)
+reset!(m) = foreach(reset!, functor(m)[1])
 
 flip(f, xs) = reverse(f.(reverse(xs)))
 
 # Vanilla RNN
 
-struct RNNCell{D,V}
-  d::D
+mutable struct RNNCell{F,A,V}
+  σ::F
+  Wi::A
+  Wh::A
+  b::V
   h::V
 end
 
-RNNCell(in::Integer, out::Integer, σ = tanh; init = initn) =
-  RNNCell(Dense(in+out, out, σ, init = init), param(init(out)))
+RNNCell(in::Integer, out::Integer, σ = tanh;
+        init = glorot_uniform) =
+  RNNCell(σ, init(out, in), init(out, out),
+          init(out), zeros(out))
 
 function (m::RNNCell)(h, x)
-  h = m.d(combine(x, h))
+  σ, Wi, Wh, b = m.σ, m.Wi, m.Wh, m.b
+  h = σ.(Wi*x .+ Wh*h .+ b)
   return h, h
 end
 
 hidden(m::RNNCell) = m.h
 
-treelike(RNNCell)
+@functor RNNCell
 
-function Base.show(io::IO, m::RNNCell)
-  print(io, "RNNCell(", m.d, ")")
+function Base.show(io::IO, l::RNNCell)
+  print(io, "RNNCell(", size(l.Wi, 2), ", ", size(l.Wi, 1))
+  l.σ == identity || print(io, ", ", l.σ)
+  print(io, ")")
 end
 
+"""
+    RNN(in::Integer, out::Integer, σ = tanh)
+
+The most basic recurrent layer; essentially acts as a `Dense` layer, but with the
+output fed back into the input each time step.
+"""
 RNN(a...; ka...) = Recur(RNNCell(a...; ka...))
 
 # LSTM
 
-struct LSTMCell{D1,D2,V}
-  forget::D1
-  input::D1
-  output::D1
-  cell::D2
-  h::V; c::V
+mutable struct LSTMCell{A,V}
+  Wi::A
+  Wh::A
+  b::V
+  h::V
+  c::V
 end
 
-function LSTMCell(in, out; init = initn)
-  cell = LSTMCell([Dense(in+out, out, σ, init = init) for _ = 1:3]...,
-                  Dense(in+out, out, tanh, init = init),
-                  param(init(out)), param(init(out)))
-  cell.forget.b.data .= 1
+function LSTMCell(in::Integer, out::Integer;
+                  init = glorot_uniform)
+  cell = LSTMCell(init(out * 4, in), init(out * 4, out), init(out * 4),
+                  zeros(out), zeros(out))
+  cell.b[gate(out, 2)] .= 1
   return cell
 end
 
-function (m::LSTMCell)(h_, x)
-  h, c = h_
-  x′ = combine(x, h)
-  forget, input, output, cell =
-    m.forget(x′), m.input(x′), m.output(x′), m.cell(x′)
+function (m::LSTMCell)((h, c), x)
+  b, o = m.b, size(h, 1)
+  g = m.Wi*x .+ m.Wh*h .+ b
+  input = σ.(gate(g, o, 1))
+  forget = σ.(gate(g, o, 2))
+  cell = tanh.(gate(g, o, 3))
+  output = σ.(gate(g, o, 4))
   c = forget .* c .+ input .* cell
-  h = output .* tanh.(c)
-  return (h, c), h
+  h′ = output .* tanh.(c)
+  return (h′, c), h′
 end
 
 hidden(m::LSTMCell) = (m.h, m.c)
 
-treelike(LSTMCell)
+@functor LSTMCell
 
-Base.show(io::IO, m::LSTMCell) =
-  print(io, "LSTMCell(",
-        size(m.forget.W, 2) - size(m.forget.W, 1), ", ",
-        size(m.forget.W, 1), ')')
+Base.show(io::IO, l::LSTMCell) =
+  print(io, "LSTMCell(", size(l.Wi, 2), ", ", size(l.Wi, 1)÷4, ")")
 
+"""
+    LSTM(in::Integer, out::Integer)
+
+[Long Short Term Memory](https://www.researchgate.net/publication/13853244_Long_Short-term_Memory)
+recurrent layer. Behaves like an RNN but generally exhibits a longer memory span over sequences.
+
+See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
+for a good overview of the internals.
+"""
 LSTM(a...; ka...) = Recur(LSTMCell(a...; ka...))
+
+# GRU
+
+mutable struct GRUCell{A,V}
+  Wi::A
+  Wh::A
+  b::V
+  h::V
+end
+
+GRUCell(in, out; init = glorot_uniform) =
+  GRUCell(init(out * 3, in), init(out * 3, out),
+          init(out * 3), zeros(out))
+
+function (m::GRUCell)(h, x)
+  b, o = m.b, size(h, 1)
+  gx, gh = m.Wi*x, m.Wh*h
+  r = σ.(gate(gx, o, 1) .+ gate(gh, o, 1) .+ gate(b, o, 1))
+  z = σ.(gate(gx, o, 2) .+ gate(gh, o, 2) .+ gate(b, o, 2))
+  h̃ = tanh.(gate(gx, o, 3) .+ r .* gate(gh, o, 3) .+ gate(b, o, 3))
+  h′ = (1 .- z).*h̃ .+ z.*h
+  return h′, h′
+end
+
+hidden(m::GRUCell) = m.h
+
+@functor GRUCell
+
+Base.show(io::IO, l::GRUCell) =
+  print(io, "GRUCell(", size(l.Wi, 2), ", ", size(l.Wi, 1)÷3, ")")
+
+"""
+    GRU(in::Integer, out::Integer)
+
+[Gated Recurrent Unit](https://arxiv.org/abs/1406.1078) layer. Behaves like an
+RNN but generally exhibits a longer memory span over sequences.
+
+See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
+for a good overview of the internals.
+"""
+GRU(a...; ka...) = Recur(GRUCell(a...; ka...))

src/layers/softmax.jl

@@ -1,23 +0,0 @@
-mutable struct Softmax{T,N,A} <: AbstractArray{T,N}
-  logits::A
-  probs::A
-  Softmax{T,N,A}(logits::A) where {T,N,A} = new(logits)
-end
-
-Softmax(logits::AbstractVecOrMat{<:AbstractFloat}) =
-  Softmax{eltype(logits),ndims(logits),typeof(logits)}(logits)
-
-@forward Softmax.logits Base.size
-
-Base.IndexStyle(::Type{Softmax{T,N,A}}) where {T,N,A} = IndexStyle(A)
-
-function Base.getindex(s::Softmax, i)
-  isdefined(s, :probs) || (s.probs = NNlib.softmax(s.logits))
-  Tracker.data(s.probs)[i]
-end
-
-softmax(xs::AbstractVecOrMat{<:AbstractFloat}) = Softmax(xs)
-
-softmax(xs::AbstractVecOrMat{<:Real}) = softmax(convert.(AbstractFloat, xs))
-
-softmax(xs::TrackedArray) = TrackedArray(Tracker.Call(NNlib.softmax, xs), Softmax(xs))

src/layers/stateless.jl

@@ -1,17 +1,296 @@
 # Cost functions
+"""
+    mae(ŷ, y)
 
-mse(ŷ, y) = sum((ŷ .- y).^2)/length(y)
+Return the mean of absolute error; calculated as
+`sum(abs.(ŷ .- y)) / length(y)`.
+"""
+mae(ŷ, y) = sum(abs.(ŷ .- y)) * 1 // length(y)
 
-crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat) =
-  -sum(y .* log.(ŷ)) / size(y, 2)
 
-@deprecate logloss(x, y) crossentropy(x, y)
+"""
+    mse(ŷ, y)
 
-function logitcrossentropy(logŷ::AbstractVecOrMat, y::AbstractVecOrMat)
-  logŷ = logŷ .- maximum(logŷ, 1)
-  ypred = logŷ .- log.(sum(exp.(logŷ), 1))
-  -sum(y .* ypred) / size(y, 2)
+Return the mean squared error between ŷ and y; calculated as
+`sum((ŷ .- y).^2) / length(y)`.
+
+# Examples
+```jldoctest
+julia> Flux.mse([0, 2], [1, 1])
+1//1
+```
+"""
+mse(ŷ, y) = sum((ŷ .- y).^2) * 1 // length(y)
+
+
+"""
+    msle(ŷ, y; ϵ=eps(eltype(ŷ)))
+
+Return the mean of the squared logarithmic errors; calculated as
+`sum((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)).^2) / length(y)`.
+The `ϵ` term provides numerical stability.
+
+Penalizes an under-predicted estimate greater than an over-predicted estimate.
+"""
+msle(ŷ, y; ϵ=eps(eltype(ŷ))) = sum((log.(ŷ .+ ϵ) .- log.(y .+ ϵ)).^2) * 1 // length(y)
+
+
+
+"""
+    huber_loss(ŷ, y; δ=1.0)
+
+Return the mean of the [Huber loss](https://en.wikipedia.org/wiki/Huber_loss)
+given the prediction `ŷ` and true values `y`.
+
+                 | 0.5 * |ŷ - y|,            for |ŷ - y| <= δ
+    Huber loss = |
+                 |  δ * (|ŷ - y| - 0.5 * δ), otherwise
+"""
+#TODO: remove dropgrad when Zygote can handle this function with CuArrays
+function huber_loss(ŷ, y;  δ=eltype(ŷ)(1))
+   abs_error = abs.(ŷ .- y)
+   temp = Zygote.dropgrad(abs_error .<  δ)
+   x = eltype(ŷ)(0.5)
+   hub_loss = sum(((abs_error.^2) .* temp) .* x .+ δ*(abs_error .- x*δ) .* (1 .- temp)) * 1 // length(y)
 end
 
-crossentropy(ŷ::Union{Softmax,TrackedArray{<:Softmax}}, y::AbstractVecOrMat) =
-  logitcrossentropy(Tracker.data(ŷ).logits, y)
+function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Nothing)
+  return -sum(xlogy.(y, ŷ)) * 1 // size(y, 2)
+end
+
+function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::Number)
+  return -sum(xlogy.(y, ŷ)) .* weight * 1 // size(y, 2)
+end
+
+function _crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat, weight::AbstractVector)
+  return -sum(xlogy.(y, ŷ) .* weight) * 1 // size(y, 2)
+end
+
+"""
+    crossentropy(ŷ, y; weight = nothing)
+
+Return the cross entropy between the given probability distributions;
+calculated as `-sum(y .* log.(ŷ) .* weight) / size(y, 2)`.
+
+`weight` can be `Nothing`, a `Number` or an `AbstractVector`.
+`weight=nothing` acts like `weight=1` but is faster.
+
+See also: [`Flux.logitcrossentropy`](@ref), [`Flux.binarycrossentropy`](@ref), [`Flux.logitbinarycrossentropy`](@ref)
+
+# Examples
+```jldoctest
+julia> Flux.crossentropy(softmax([-1.1491, 0.8619, 0.3127]), [1, 1, 0])
+3.085467254747739
+```
+"""
+crossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight=nothing) = _crossentropy(ŷ, y, weight)
+
+"""
+    logitcrossentropy(ŷ, y; weight = 1)
+
+Return the crossentropy computed after a [`Flux.logsoftmax`](@ref) operation;
+calculated as `-sum(y .* logsoftmax(ŷ) .* weight) / size(y, 2)`.
+
+`logitcrossentropy(ŷ, y)` is mathematically equivalent to
+[`Flux.crossentropy(softmax(ŷ), y)`](@ref) but it is more numerically stable.
+
+See also: [`Flux.crossentropy`](@ref), [`Flux.binarycrossentropy`](@ref), [`Flux.logitbinarycrossentropy`](@ref)
+
+# Examples
+```jldoctest
+julia> Flux.logitcrossentropy([-1.1491, 0.8619, 0.3127], [1, 1, 0])
+3.085467254747738
+```
+"""
+function logitcrossentropy(ŷ::AbstractVecOrMat, y::AbstractVecOrMat; weight = 1)
+  return -sum(y .* logsoftmax(ŷ) .* weight) * 1 // size(y, 2)
+end
+
+"""
+    binarycrossentropy(ŷ, y; ϵ=eps(ŷ))
+
+Return ``-y*\\log(ŷ + ϵ) - (1-y)*\\log(1-ŷ + ϵ)``. The `ϵ` term provides numerical stability.
+
+Typically, the prediction `ŷ` is given by the output of a [`sigmoid`](@ref) activation.
+
+See also: [`Flux.crossentropy`](@ref), [`Flux.logitcrossentropy`](@ref), [`Flux.logitbinarycrossentropy`](@ref)
+
+# Examples
+```jldoctest
+julia> Flux.binarycrossentropy.(σ.([-1.1491, 0.8619, 0.3127]), [1, 1, 0])
+3-element Array{Float64,1}:
+ 1.424397097347566
+ 0.35231664672364077
+ 0.8616703662235441
+```
+"""
+binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -xlogy(y, ŷ + ϵ) - xlogy(1 - y, 1 - ŷ + ϵ)
+
+# Re-definition to fix interaction with CuArrays.
+CuArrays.@cufunc binarycrossentropy(ŷ, y; ϵ=eps(ŷ)) = -y*log(ŷ + ϵ) - (1 - y)*log(1 - ŷ + ϵ)
+
+"""
+    logitbinarycrossentropy(ŷ, y)
+
+`logitbinarycrossentropy(ŷ, y)` is mathematically equivalent to
+[`Flux.binarycrossentropy(σ(ŷ), y)`](@ref) but it is more numerically stable.
+
+See also: [`Flux.crossentropy`](@ref), [`Flux.logitcrossentropy`](@ref), [`Flux.binarycrossentropy`](@ref)
+
+# Examples
+```jldoctest
+julia> Flux.logitbinarycrossentropy.([-1.1491, 0.8619, 0.3127], [1, 1, 0])
+3-element Array{Float64,1}:
+ 1.4243970973475661
+ 0.35231664672364094
+ 0.8616703662235443
+```
+"""
+logitbinarycrossentropy(ŷ, y) = (1 - y)*ŷ - logσ(ŷ)
+
+# Re-definition to fix interaction with CuArrays.
+CuArrays.@cufunc logitbinarycrossentropy(ŷ, y) = (1 - y)*ŷ - logσ(ŷ)
+
+"""
+    normalise(x; dims=1)
+
+Normalise `x` to mean 0 and standard deviation 1 across the dimensions given by `dims`.
+Defaults to normalising over columns.
+
+```jldoctest
+julia> a = reshape(collect(1:9), 3, 3)
+3×3 Array{Int64,2}:
+ 1  4  7
+ 2  5  8
+ 3  6  9
+
+julia> Flux.normalise(a)
+3×3 Array{Float64,2}:
+ -1.22474  -1.22474  -1.22474
+  0.0       0.0       0.0
+  1.22474   1.22474   1.22474
+
+julia> Flux.normalise(a, dims=2)
+3×3 Array{Float64,2}:
+ -1.22474  0.0  1.22474
+ -1.22474  0.0  1.22474
+ -1.22474  0.0  1.22474
+```
+"""
+function normalise(x::AbstractArray; dims=1)
+  μ′ = mean(x, dims = dims)
+  σ′ = std(x, dims = dims, mean = μ′, corrected=false)
+  return (x .- μ′) ./ σ′
+end
+
+"""
+    kldivergence(ŷ, y)
+
+Return the
+[Kullback-Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence)
+between the given probability distributions.
+
+KL divergence is a measure of how much one probability distribution is different
+from the other.
+It is always non-negative and zero only when both the distributions are equal
+everywhere.
+"""
+function kldivergence(ŷ, y)
+  entropy = sum(xlogx.(y)) * 1 //size(y,2)
+  cross_entropy = crossentropy(ŷ, y)
+  return entropy + cross_entropy
+end
+
+"""
+    poisson(ŷ, y)
+
+Return how much the predicted distribution `ŷ` diverges from the expected Poisson
+distribution `y`; calculated as `sum(ŷ .- y .* log.(ŷ)) / size(y, 2)`.
+
+[More information.](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/poisson).
+"""
+poisson(ŷ, y) = sum(ŷ .- xlogy.(y, ŷ)) * 1 // size(y,2)
+
+"""
+    hinge(ŷ, y)
+
+Return the [hinge loss](https://en.wikipedia.org/wiki/Hinge_loss) given the
+prediction `ŷ` and true labels `y` (containing 1 or -1); calculated as
+`sum(max.(0, 1 .- ŷ .* y)) / size(y, 2)`.
+
+See also: [`squared_hinge`](@ref)
+"""
+hinge(ŷ, y) = sum(max.(0, 1 .-  ŷ .* y)) * 1 // size(y, 2)
+
+"""
+    squared_hinge(ŷ, y)
+
+Return the squared hinge loss given the prediction `ŷ` and true labels `y`
+(containing 1 or -1); calculated as `sum((max.(0, 1 .- ŷ .* y)).^2) / size(y, 2)`.
+
+See also: [`hinge`](@ref)
+"""
+squared_hinge(ŷ, y) = sum((max.(0, 1 .- ŷ .* y)).^2) * 1 // size(y, 2)
+
+"""
+    dice_coeff_loss(ŷ, y; smooth=1)
+
+Return a loss based on the dice coefficient.
+Used in the [V-Net](https://arxiv.org/pdf/1606.04797v1.pdf) image segmentation
+architecture.
+Similar to the F1_score. Calculated as:
+    1 - 2*sum(|ŷ .* y| + smooth) / (sum(ŷ.^2) + sum(y.^2) + smooth)`
+"""
+dice_coeff_loss(ŷ, y; smooth=eltype(ŷ)(1.0)) = 1 - (2*sum(y .* ŷ) + smooth) / (sum(y.^2) + sum(ŷ.^2) + smooth)
+
+"""
+    tversky_loss(ŷ, y; β=0.7)
+
+Return the [Tversky loss](https://arxiv.org/pdf/1706.05721.pdf).
+Used with imbalanced data to give more weight to false negatives.
+Larger β weigh recall higher than precision (by placing more emphasis on false negatives)
+Calculated as:
+    1 - sum(|y .* ŷ| + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)
+"""
+tversky_loss(ŷ, y; β=eltype(ŷ)(0.7)) = 1 - (sum(y .* ŷ) + 1) / (sum(y .* ŷ + β*(1 .- y) .* ŷ + (1 - β)*y .* (1 .- ŷ)) + 1)
+
+"""
+    flatten(x::AbstractArray)
+
+Transform (w, h, c, b)-shaped input into (w × h × c, b)-shaped output
+by linearizing all values for each element in the batch.
+"""
+function flatten(x::AbstractArray)
+  return reshape(x, :, size(x)[end])
+end
+
+"""
+    xlogx(x)
+Return `x * log(x)` for `x ≥ 0`, handling `x = 0` by taking the downward limit.
+"""
+function xlogx(x)
+  result = x * log(x)
+  ifelse(iszero(x), zero(result), result)
+end
+CuArrays.@cufunc function xlogx(x)
+  result = x * log(x)
+  ifelse(iszero(x), zero(result), result)
+end
+
+"""
+    xlogy(x, y)
+Return `x * log(y)` for `y > 0` with correct limit at `x = 0`.
+"""
+function xlogy(x, y)
+  result = x * log(y)
+  ifelse(iszero(x), zero(result), result)
+end
+CuArrays.@cufunc function xlogy(x, y)
+  result = x * log(y)
+  ifelse(iszero(x), zero(result), result)
+end
+
+@adjoint function broadcasted(::typeof(xlogy), x::Zygote.Numeric, y::Zygote.Numeric)
+  res = xlogy.(x, y)
+  res, Δ -> (nothing, Zygote.unbroadcast(x, xlogy.(Δ, y)), Zygote.unbroadcast(y, Δ .* x ./ y))
+end

src/onehot.jl

@@ -1,3 +1,5 @@
+import Base: *
+
 struct OneHotVector <: AbstractVector{Bool}
   ix::UInt32
   of::UInt32
@@ -7,7 +9,9 @@ Base.size(xs::OneHotVector) = (Int64(xs.of),)
 
 Base.getindex(xs::OneHotVector, i::Integer) = i == xs.ix
 
-Base.:*(A::AbstractMatrix, b::OneHotVector) = A[:, b.ix]
+Base.getindex(xs::OneHotVector, ::Colon) = OneHotVector(xs.ix, xs.of)
+
+A::AbstractMatrix * b::OneHotVector = A[:, b.ix]
 
 struct OneHotMatrix{A<:AbstractVector{OneHotVector}} <: AbstractMatrix{Bool}
   height::Int
@@ -16,34 +20,106 @@ end
 
 Base.size(xs::OneHotMatrix) = (Int64(xs.height),length(xs.data))
 
-Base.getindex(xs::OneHotMatrix, i::Int, j::Int) = xs.data[j][i]
+Base.getindex(xs::OneHotMatrix, i::Union{Integer, AbstractVector}, j::Integer) = xs.data[j][i]
+Base.getindex(xs::OneHotMatrix, ::Colon, i::Integer) = xs.data[i]
+Base.getindex(xs::OneHotMatrix, ::Colon, i::AbstractArray) = OneHotMatrix(xs.height, xs.data[i])
+Base.getindex(xs::OneHotMatrix, ::Colon, ::Colon) = OneHotMatrix(xs.height, copy(xs.data))
 
-Base.:*(A::AbstractMatrix, B::OneHotMatrix) = A[:, map(x->x.ix, B.data)]
+Base.getindex(xs::OneHotMatrix, i::Integer, ::Colon) = map(x -> x[i], xs.data)
+
+# remove workaround when https://github.com/JuliaGPU/CuArrays.jl/issues/676 is fixed
+A::AbstractMatrix * B::OneHotMatrix = A[:, cpu(map(x->x.ix, B.data))]
 
 Base.hcat(x::OneHotVector, xs::OneHotVector...) = OneHotMatrix(length(x), [x, xs...])
 
 batch(xs::AbstractArray{<:OneHotVector}) = OneHotMatrix(length(first(xs)), xs)
 
-import NNlib.adapt
+import Adapt: adapt, adapt_structure
 
-adapt(T, xs::OneHotMatrix) = OneHotMatrix(xs.height, adapt(T, xs.data))
+adapt_structure(T, xs::OneHotMatrix) = OneHotMatrix(xs.height, adapt(T, xs.data))
 
-@require CuArrays begin
-  import CuArrays: CuArray, cudaconvert
-  Base.Broadcast._containertype(::Type{<:OneHotMatrix{<:CuArray}}) = CuArray
-  cudaconvert(x::OneHotMatrix{<:CuArray}) = OneHotMatrix(x.height, cudaconvert(x.data))
-end
+import .CuArrays: CuArray, CuArrayStyle, cudaconvert
+import Base.Broadcast: BroadcastStyle, ArrayStyle
+BroadcastStyle(::Type{<:OneHotMatrix{<:CuArray}}) = CuArrayStyle{2}()
+cudaconvert(x::OneHotMatrix{<:CuArray}) = OneHotMatrix(x.height, cudaconvert(x.data))
 
+"""
+    onehot(l, labels[, unk])
+
+Create a `OneHotVector` with its `l`-th element `true` based on the
+possible set of `labels`.
+If `unk` is given, return `onehot(unk, labels)` if the input label `l` is not found
+in `labels`; otherwise, it will raise an error.
+
+# Examples
+```jldoctest
+julia> Flux.onehot(:b, [:a, :b, :c])
+3-element Flux.OneHotVector:
+ 0
+ 1
+ 0
+
+julia> Flux.onehot(:c, [:a, :b, :c])
+3-element Flux.OneHotVector:
+ 0
+ 0
+ 1
+```
+"""
 function onehot(l, labels)
-  i = findfirst(labels, l)
+  i = something(findfirst(isequal(l), labels), 0)
   i > 0 || error("Value $l is not in labels")
   OneHotVector(i, length(labels))
 end
 
-onehotbatch(ls, labels) = OneHotMatrix(length(labels), [onehot(l, labels) for l in ls])
+function onehot(l, labels, unk)
+  i = something(findfirst(isequal(l), labels), 0)
+  i > 0 || return onehot(unk, labels)
+  OneHotVector(i, length(labels))
+end
 
-argmax(y::AbstractVector, labels = 1:length(y)) =
-  labels[findfirst(y, maximum(y))]
+"""
+    onehotbatch(ls, labels[, unk...])
 
-argmax(y::AbstractMatrix, l...) =
-  squeeze(mapslices(y -> argmax(y, l...), y, 1), 1)
+Create a `OneHotMatrix` with a batch of labels based on the
+possible set of `labels`.
+If `unk` is given, return [`onehot(unk, labels)`](@ref) if one of the input
+labels `ls` is not found in `labels`; otherwise it will error.
+
+# Examples
+```jldoctest
+julia> Flux.onehotbatch([:b, :a, :b], [:a, :b, :c])
+3×3 Flux.OneHotMatrix{Array{Flux.OneHotVector,1}}:
+ 0  1  0
+ 1  0  1
+ 0  0  0
+```
+"""
+onehotbatch(ls, labels, unk...) =
+  OneHotMatrix(length(labels), [onehot(l, labels, unk...) for l in ls])
+
+Base.argmax(xs::OneHotVector) = xs.ix
+
+"""
+    onecold(y[, labels = 1:length(y)])
+
+Inverse operations of [`onehot`](@ref).
+
+# Examples
+```jldoctest
+julia> Flux.onecold([true, false, false], [:a, :b, :c])
+:a
+
+julia> Flux.onecold([0.3, 0.2, 0.5], [:a, :b, :c])
+:c
+```
+"""
+onecold(y::AbstractVector, labels = 1:length(y)) = labels[Base.argmax(y)]
+
+onecold(y::AbstractMatrix, labels...) =
+  dropdims(mapslices(y -> onecold(y, labels...), y, dims=1), dims=1)
+
+onecold(y::OneHotMatrix, labels...) =
+  mapreduce(x -> Flux.onecold(x, labels...), |, y.data, dims = 2, init = 0)
+
+@nograd onecold, onehot, onehotbatch

src/optimise/Optimise.jl

@@ -1,21 +1,14 @@
 module Optimise
 
-export update!, params, train!,
-  SGD, ADAM, Momentum, Nesterov, RMSProp, ADAGrad, ADADelta
+using LinearAlgebra
 
-struct Param{T}
-  x::T
-  Δ::T
-end
-
-Base.convert(::Type{Param}, x::AbstractArray) = Param(x, zeros(x))
+export train!, update!,
+	Descent, ADAM, Momentum, Nesterov, RMSProp,
+	ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW,RADAM, 
+	InvDecay, ExpDecay, WeightDecay, stop, Optimiser,
+	ClipValue, ClipNorm
 
 include("optimisers.jl")
-include("interface.jl")
 include("train.jl")
 
-using Flux.Tracker: TrackedArray
-
-Base.convert(::Type{Param}, x::TrackedArray) = Param(x.data, x.grad[])
-
 end

src/optimise/interface.jl

@@ -1,18 +0,0 @@
-call(f, xs...) = f(xs...)
-
-function optimiser(ps, fs...)
-  ps = [Param(p) for p in ps]
-  fs = map(ps) do p
-    os = map(f -> f(p), fs)
-    () -> foreach(call, os)
-  end
-  () -> foreach(call, fs)
-end
-
-SGD(ps, η = 1) = optimiser(ps, p -> descent(p, η))
-ADAM(ps, η = 0.001, β1 = 0.9, β2 = 0.999, ϵ = 1e-08, decay = 0.0) = optimiser(ps, p -> adam(p; η = η, β1 = β1, β2 = β2, ϵ = ϵ), p -> invdecay(p, decay), p -> descent(p, 1))
-Momentum(ps,ρ, decay = 0.0) = optimiser(ps, p -> momentum(p, ρ), p -> invdecay(p, decay), p -> descent(p, 1))
-Nesterov(ps,ρ, decay = 0.0) = optimiser(ps, p -> nesterov(p, ρ), p -> invdecay(p, decay), p -> descent(p, 1))
-RMSProp(ps, η = 0.001, ρ = 0.9, ϵ = 1e-8, decay = 0.0) = optimiser(ps, p -> rmsprop(p; η = η, ρ = ρ, ϵ = ϵ), p -> invdecay(p, decay), p -> descent(p, 1))
-ADAGrad(ps, η = 0.01, ϵ = 1e-8, decay = 0.0) = optimiser(ps, p -> adagrad(p; η = η, ϵ = ϵ), p -> invdecay(p, decay), p -> descent(p, 1))
-ADADelta(ps, η = 0.01, ρ = 0.95, ϵ = 1e-8, decay = 0.0) = optimiser(ps, p -> adadelta(p; ρ = ρ, ϵ = ϵ), p -> invdecay(p, decay), p -> descent(p, 1))

src/optimise/optimisers.jl

@@ -1,74 +1,563 @@
-function descent(p::Param, η::Real)
-  function ()
-    p.x .-= p.Δ .* η
-    p.Δ .= 0
+using Flux
+using MacroTools: @forward
+
+const ϵ = 1e-8
+
+# TODO: should use weak refs
+
+"""
+    Descent(η = 0.1)
+
+Classic gradient descent optimiser with learning rate `η`.
+For each parameter `p` and its gradient `δp`, this runs `p -= η*δp`
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+
+# Examples
+```julia
+opt = Descent()
+
+opt = Descent(0.3)
+
+ps = params(model)
+
+gs = gradient(ps) do
+    loss(x, y)
+end
+
+Flux.Optimise.update!(opt, ps, gs)
+```
+"""
+mutable struct Descent
+  eta::Float64
+end
+
+Descent() = Descent(0.1)
+
+function apply!(o::Descent, x, Δ)
+  Δ .*= o.eta
+end
+
+"""
+    Momentum(η = 0.01, ρ = 0.9)
+
+Gradient descent optimizer with learning rate `η` and momentum `ρ`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Momentum (`ρ`): Controls the acceleration of gradient descent in the
+                  prominent direction, in effect dampening oscillations.
+
+# Examples
+```julia
+opt = Momentum()
+
+opt = Momentum(0.01, 0.99)
+```
+"""
+mutable struct Momentum
+  eta::Float64
+  rho::Float64
+  velocity::IdDict
+end
+
+Momentum(η = 0.01, ρ = 0.9) = Momentum(η, ρ, IdDict())
+
+function apply!(o::Momentum, x, Δ)
+  η, ρ = o.eta, o.rho
+  v = get!(o.velocity, x, zero(x))::typeof(x)
+  @. v = ρ * v - η * Δ
+  @. Δ = -v
+end
+
+"""
+    Nesterov(η = 0.001, ρ = 0.9)
+
+Gradient descent optimizer with learning rate `η` and Nesterov momentum `ρ`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Nesterov momentum (`ρ`): Controls the acceleration of gradient descent in the
+                           prominent direction, in effect dampening oscillations.
+
+# Examples
+```julia
+opt = Nesterov()
+
+opt = Nesterov(0.003, 0.95)
+```
+"""
+mutable struct Nesterov
+  eta::Float64
+  rho::Float64
+  velocity::IdDict
+end
+
+Nesterov(η = 0.001, ρ = 0.9) = Nesterov(η, ρ, IdDict())
+
+function apply!(o::Nesterov, x, Δ)
+  η, ρ = o.eta, o.rho
+  v = get!(o.velocity, x, zero(x))::typeof(x)
+  d = @. ρ^2 * v - (1+ρ) * η * Δ
+  @. v = ρ*v - η*Δ
+  @. Δ = -d
+end
+
+"""
+    RMSProp(η = 0.001, ρ = 0.9)
+
+Optimizer using the
+[RMSProp](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
+algorithm. Often a good choice for recurrent networks. Parameters other than learning rate
+generally don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Momentum (`ρ`): Controls the acceleration of gradient descent in the
+                  prominent direction, in effect dampening oscillations.
+
+# Examples
+```julia
+opt = RMSProp()
+
+opt = RMSProp(0.002, 0.95)
+```
+"""
+mutable struct RMSProp
+  eta::Float64
+  rho::Float64
+  acc::IdDict
+end
+
+RMSProp(η = 0.001, ρ = 0.9) = RMSProp(η, ρ, IdDict())
+
+function apply!(o::RMSProp, x, Δ)
+  η, ρ = o.eta, o.rho
+  acc = get!(o.acc, x, zero(x))::typeof(x)
+  @. acc = ρ * acc + (1 - ρ) * Δ^2
+  @. Δ *= η / (√acc + ϵ)
+end
+
+"""
+    ADAM(η = 0.001, β::Tuple = (0.9, 0.999))
+
+[ADAM](https://arxiv.org/abs/1412.6980v8) optimiser.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+
+# Examples
+```julia
+opt = ADAM()
+
+opt = ADAM(0.001, (0.9, 0.8))
+```
+"""
+mutable struct ADAM
+  eta::Float64
+  beta::Tuple{Float64,Float64}
+  state::IdDict
+end
+
+ADAM(η = 0.001, β = (0.9, 0.999)) = ADAM(η, β, IdDict())
+
+function apply!(o::ADAM, x, Δ)
+  η, β = o.eta, o.beta
+  mt, vt, βp = get!(o.state, x, (zero(x), zero(x), β))
+  @. mt = β[1] * mt + (1 - β[1]) * Δ
+  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
+  @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η
+  o.state[x] = (mt, vt, βp .* β)
+  return Δ
+end
+
+"""
+    RADAM(η = 0.001, β::Tuple = (0.9, 0.999))
+
+[Rectified ADAM](https://arxiv.org/pdf/1908.03265v1.pdf) optimizer.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+
+# Examples
+```julia
+opt = RADAM()
+
+opt = RADAM(0.001, (0.9, 0.8))
+```
+"""
+mutable struct RADAM
+  eta::Float64
+  beta::Tuple{Float64,Float64}
+  state::IdDict
+end
+
+RADAM(η = 0.001, β = (0.9, 0.999)) = RADAM(η, β, IdDict())
+
+function apply!(o::RADAM, x, Δ)
+  η, β = o.eta, o.beta
+  ρ∞ = 2/(1-β[2])-1
+  mt, vt, βp, t = get!(o.state, x, (zero(x), zero(x), β, 1))
+  @. mt = β[1] * mt + (1 - β[1]) * Δ
+  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
+  ρ = ρ∞ - 2t*βp[2]/(1-βp[2])
+  if ρ > 4
+    r = sqrt((ρ-4)*(ρ-2)*ρ∞/((ρ∞-4)*(ρ∞-2)*ρ))
+    @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η * r
+  else
+    @. Δ =  mt / (1 - βp[1]) * η
   end
+  o.state[x] = (mt, vt, βp .* β, t+1)
+  return Δ
 end
 
-function momentum(p::Param, ρ::Real)
-  mo = zeros(p.x)
-  () -> p.Δ .= mo .= ρ .* mo .+ p.Δ
+"""
+    AdaMax(η = 0.001, β::Tuple = (0.9, 0.999))
+
+[AdaMax](https://arxiv.org/abs/1412.6980v9) is a variant of ADAM based on the ∞-norm.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+
+# Examples
+```julia
+opt = AdaMax()
+
+opt = AdaMax(0.001, (0.9, 0.995))
+```
+"""
+mutable struct AdaMax
+  eta::Float64
+  beta::Tuple{Float64,Float64}
+  state::IdDict
 end
 
-function nesterov(p::Param, ρ::Real)
-  mo = zeros(p.x)
-  function ()
-    mo  .= ρ .* mo .+ p.Δ
-    p.Δ .= ρ .* mo .+ p.Δ
+AdaMax(η = 0.001, β = (0.9, 0.999)) = AdaMax(η, β, IdDict())
+
+function apply!(o::AdaMax, x, Δ)
+  η, β = o.eta, o.beta
+  mt, ut, βp = get!(o.state, x, (zero(x), zero(x), β))
+  @. mt = β[1] * mt + (1 - β[1]) * Δ
+  @. ut = max(β[2] * ut, abs(Δ))
+  @. Δ = (η/(1 - βp[1])) * mt/(ut + ϵ)
+  o.state[x] = (mt, ut, βp .* β)
+  return Δ
+end
+
+"""
+    ADAGrad(η = 0.1)
+
+[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimizer. It has
+parameter specific learning rates based on how frequently it is updated.
+Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+
+# Examples
+```julia
+opt = ADAGrad()
+
+opt = ADAGrad(0.001)
+```
+"""
+mutable struct ADAGrad
+  eta::Float64
+  acc::IdDict
+end
+
+ADAGrad(η = 0.1) = ADAGrad(η, IdDict())
+
+function apply!(o::ADAGrad, x, Δ)
+  η = o.eta
+  acc = get!(o.acc, x, fill!(zero(x), ϵ))::typeof(x)
+  @. acc += Δ^2
+  @. Δ *= η / (√acc + ϵ)
+end
+
+"""
+    ADADelta(ρ = 0.9)
+
+[ADADelta](https://arxiv.org/abs/1212.5701) is a version of ADAGrad adapting its learning
+rate based on a window of past gradient updates.
+Parameters don't need tuning.
+
+# Parameters
+- Rho (`ρ`): Factor by which the gradient is decayed at each time step.
+
+# Examples
+```julia
+opt = ADADelta()
+
+opt = ADADelta(0.89)
+```
+"""
+mutable struct ADADelta
+  rho::Float64
+  state::IdDict
+end
+
+ADADelta(ρ = 0.9) = ADADelta(ρ, IdDict())
+
+function apply!(o::ADADelta, x, Δ)
+  ρ = o.rho
+  acc, Δacc = get!(o.state, x, (zero(x), zero(x)))
+  @. acc = ρ * acc + (1 - ρ) * Δ^2
+  @. Δ *= √Δacc/ (√acc + ϵ)
+  @. Δacc = ρ * Δacc + (1 - ρ) * Δ^2
+  return Δ
+end
+
+"""
+    AMSGrad(η = 0.001, β::Tuple = (0.9, 0.999))
+
+The [AMSGrad](https://openreview.net/forum?id=ryQu7f-RZ) version of the ADAM
+optimiser. Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+
+# Examples
+```julia
+opt = AMSGrad()
+
+opt = AMSGrad(0.001, (0.89, 0.995))
+```
+"""
+mutable struct AMSGrad
+  eta::Float64
+  beta::Tuple{Float64, Float64}
+  state::IdDict
+end
+
+AMSGrad(η = 0.001, β = (0.9, 0.999)) = AMSGrad(η, β, IdDict())
+
+function apply!(o::AMSGrad, x, Δ)
+  η, β = o.eta, o.beta
+  mt, vt, v̂t = get!(o.state, x, (fill!(zero(x), ϵ), fill!(zero(x), ϵ), fill!(zero(x), ϵ)))
+  @. mt = β[1] * mt + (1 - β[1]) * Δ
+  @. vt = β[2] * vt + (1 - β[2]) * Δ ^ 2
+  @. v̂t = max(v̂t, vt)
+  @. Δ = η * mt / (√v̂t + ϵ)
+end
+
+"""
+    NADAM(η = 0.001, β::Tuple = (0.9, 0.999))
+
+[NADAM](http://cs229.stanford.edu/proj2015/054_report.pdf) is a Nesterov variant of ADAM.
+Parameters don't need tuning.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+
+# Examples
+```julia
+opt = NADAM()
+
+opt = NADAM(0.002, (0.89, 0.995))
+```
+"""
+mutable struct NADAM
+  eta::Float64
+  beta::Tuple{Float64, Float64}
+  state::IdDict
+end
+
+NADAM(η = 0.001, β = (0.9, 0.999)) = NADAM(η, β, IdDict())
+
+function apply!(o::NADAM, x, Δ)
+  η, β = o.eta, o.beta
+  mt, vt, (β1p, β2p) = get!(o.state, x, (zero(x), zero(x), o.beta))
+  @. mt = β[1] * mt + (1 - β[1]) * Δ
+  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
+  @. Δ = (β[1] * mt / (1 - β[1] * β1p) + (1 - β[1]) * Δ / (1 - β1p)) / (√(vt * β[2] / (1 - β2p)) + ϵ) * η
+  o.state[x] = (mt, vt, (β1p * β[1], β2p * β[2]))
+  return Δ
+end
+
+"""
+    ADAMW(η = 0.001, β::Tuple = (0.9, 0.999), decay = 0)
+
+[ADAMW](https://arxiv.org/abs/1711.05101) is a variant of ADAM fixing (as in repairing) its
+weight decay regularization.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- Decay of momentums (`β::Tuple`): Exponential decay for the first (β1) and the
+                                   second (β2) momentum estimate.
+- `decay`: Decay applied to weights during optimisation.
+
+# Examples
+```julia
+opt = ADAMW()
+
+opt = ADAMW(0.001, (0.89, 0.995), 0.1)
+```
+"""
+ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) =
+  Optimiser(ADAM(η, β), WeightDecay(decay))
+
+# Compose optimizers
+
+"""
+    Optimiser(a, b, c...)
+
+Combine several optimisers into one; each optimiser produces a modified gradient
+that will be fed into the next, and this is finally applied to the parameter as
+usual.
+"""
+mutable struct Optimiser
+  os::Vector{Any}
+end
+
+Optimiser(o...) = Optimiser(Any[o...])
+
+@forward Optimiser.os Base.getindex, Base.first, Base.last, Base.lastindex, Base.push!, Base.setindex!
+@forward Optimiser.os Base.iterate
+
+Base.getindex(c::Optimiser, i::AbstractArray) = Optimiser(c.os[i]...)
+
+function apply!(o::Optimiser, x, Δ)
+  for opt in o.os
+    Δ = apply!(opt, x, Δ)
   end
+  return Δ
 end
 
-function clip(p::Param, thresh::Real)
-  () -> clamp!(p.Δ, -thresh, thresh)
+"""
+    InvDecay(γ = 0.001)
+
+Apply inverse time decay to an optimiser, so that the effective step size at
+iteration `n` is `eta / (1 + γ * n)` where `eta` is the initial step size.
+The wrapped optimiser's step size is not modified.
+
+# Examples
+```julia
+Optimiser(InvDecay(..), Opt(..))
+```
+"""
+mutable struct InvDecay
+  gamma::Float64
+  state::IdDict
 end
 
-function weightdecay(p::Param, γ::Real)
-  () -> p.Δ .+= γ .* p.x
+InvDecay(γ = 0.001) = InvDecay(γ, IdDict())
+
+function apply!(o::InvDecay, x, Δ)
+  γ = o.gamma
+  n = get!(o.state, x, 1)
+  Δ .*= 1 / (1 + γ * n)
+  o.state[x] = n + 1
+  return Δ
 end
 
-function invdecay(p::Param, γ::Real)
-  n = 0
-  function ()
-    p.Δ .*= 1 / (1 + γ * n)
-    n += 1
+"""
+    ExpDecay(η = 0.001, decay = 0.1, decay_step = 1000, clip = 1e-4)
+
+Discount the learning rate `η` by the factor `decay` every `decay_step` steps till
+a minimum of `clip`.
+
+# Parameters
+- Learning rate (`η`): Amount by which gradients are discounted before updating
+                       the weights.
+- `decay`: Factor by which the learning rate is discounted.
+- `decay_step`: Schedule decay operations by setting the number of steps between
+                two decay operations.
+- `clip`: Minimum value of learning rate.
+
+# Examples
+To apply exponential decay to an optimiser:
+```julia
+Optimiser(ExpDecay(..), Opt(..))
+
+opt = Optimiser(ExpDecay(), ADAM())
+```
+"""
+mutable struct ExpDecay
+  eta::Float64
+  decay::Float64
+  step::Int64
+  clip::Float64
+  current::IdDict
+end
+
+ExpDecay(opt = 0.001, decay = 0.1, decay_step = 1000, clip = 1e-4) = ExpDecay(opt, decay, decay_step, clip, IdDict())
+
+function apply!(o::ExpDecay, x, Δ)
+  η, s, decay = o.eta, o.step, o.decay
+  n = o.current[x] = get(o.current, x, 0) + 1
+  if o.current[x]%s == 0 && count(x -> x%s == 0, values(o.current)) == 1
+    η = max(η * decay, o.clip)
+    o.eta = η
   end
+  @. Δ *= η
 end
 
-function rmsprop(p::Param; η::Real = 0.001, ρ::Real = 0.9, ϵ::Real = 1e-8)
-  acc  = zeros(p.x) .+ ϵ
-  function ()
-    @. acc = ρ * acc + (1 - ρ) * p.Δ ^ 2
-    @. p.Δ /= √acc * η
-  end
+"""
+    WeightDecay(wd = 0)
+
+Decay weights by `wd`.
+
+# Parameters
+- Weight decay (`wd`)
+"""
+mutable struct WeightDecay
+  wd::Real
 end
 
-function adagrad(p::Param; η::Real = 0.01, ϵ::Real = 1e-8)
-  acc = zeros(p.x) .+ ϵ
-  function ()
-    @. acc += p.Δ ^ 2
-    @. p.Δ /= √acc * η
-  end
+WeightDecay() = WeightDecay(0)
+
+function apply!(o::WeightDecay, x, Δ)
+  wd = o.wd
+  @. Δ += wd * x
 end
 
-function adadelta(p::Param; ρ::Real = 0.95, ϵ::Real = 1e-8)
-  acc = zeros(p.x) .+ ϵ
-  Δacc = zeros(p.x) .+ ϵ
-  function ()
-    @. acc = ρ * acc + (1 - ρ) * p.Δ ^ 2
-    @. p.Δ *= √Δacc / √acc
-    @. Δacc = ρ * Δacc + (1 - ρ) * p.Δ ^ 2
-  end
+"""
+    ClipValue(thresh)
+
+Clip gradients when their absolute value exceeds `thresh`.
+"""
+mutable struct ClipValue{T}
+    thresh::T
 end
 
-function adam(p::Param; η::Real = 0.001, β1::Real = 0.9, β2::Real = 0.999, ϵ::Real = 1e-8)
-  mt = zeros(p.x)
-  vt = zeros(p.x) .+ ϵ
-  β1p, β2p = β1, β2
-  function ()
-    @. mt = β1 * mt + (1 - β1) * p.Δ
-    @. vt = β2 * vt + (1 - β2) * p.Δ ^ 2
-    @. p.Δ = √(1 - β2p) / √(1 - β1p) * mt / √vt * η
-    β1p *= β1
-    β2p *= β2
-  end
+apply!(o::ClipValue, x, Δ) = clamp!(Δ, -o.thresh, o.thresh)
+
+"""
+    ClipNorm(thresh)
+
+Clip gradients when their L2 norm exceeds `thresh`.
+"""
+mutable struct ClipNorm{T}
+    thresh::T
 end
+
+function apply!(o::ClipNorm, x, Δ)
+    Δnrm = norm(Δ)
+    if Δnrm > o.thresh
+        rmul!(Δ, o.thresh / Δnrm)
+    end
+    return Δ
+end

src/optimise/train.jl

@@ -1,26 +1,123 @@
 using Juno
-using Flux.Tracker: back!
+import Zygote: Params, gradient
+
 
-tocb(f) = f
-tocb(fs::AbstractVector) = () -> foreach(call, fs)
 
 """
-    train!(loss, data, opt; cb = () -> ())
+    update!(x, x̄)
 
-For each datapoint `d` in `data` computes the gradient of `loss(d...)` through
-backpropagation and calls the optimizer `opt` and the callback `cb`
-(i.e. `opt()` and `cb()`).
-
-Multiple callbacks can be passed to `cb` as an array.
+Update the array `x` according to `x .-= x̄`.
 """
-function train!(loss, data, opt; cb = () -> ())
-  cb = tocb(cb)
-  @progress for d in data
-    l = loss(d...)
-    isinf(l.data[]) && error("Loss is Inf")
-    isnan(l.data[]) && error("Loss is NaN")
-    back!(l)
-    opt()
-    cb()
+function update!(x::AbstractArray, x̄)
+  x .-= x̄
+end
+
+"""
+    update!(opt, p, g)
+    update!(opt, ps::Params, gs)
+
+Perform an update step of the parameters `ps` (or the single parameter `p`)
+according to optimizer `opt`  and the gradients `gs` (the gradient `g`).
+
+As a result, the parameters are mutated and the optimizer's internal state may change.
+"""
+function update!(opt, x, x̄)
+  x .-= apply!(opt, x, x̄)
+end
+
+function update!(opt, xs::Params, gs)
+  for x in xs
+    gs[x] == nothing && continue
+    update!(opt, x, gs[x])
   end
 end
+
+# Callback niceties
+call(f, xs...) = f(xs...)
+runall(f) = f
+runall(fs::AbstractVector) = () -> foreach(call, fs)
+
+struct StopException <: Exception end
+
+"""
+    stop()
+
+Call `Flux.stop()` in a callback to indicate when a callback condition is met.
+This will trigger the train loop to stop and exit.
+
+# Examples
+```julia
+cb = function ()
+  accuracy() > 0.9 && Flux.stop()
+end
+```
+"""
+function stop()
+  throw(StopException())
+end
+
+"""
+    train!(loss, params, data, opt; cb)
+
+For each datapoint `d` in `data` compute the gradient of `loss(d...)` through
+backpropagation and call the optimizer `opt`.
+
+In case datapoints `d` are of numeric array type, assume no splatting is needed
+and compute the gradient of `loss(d)`.
+
+A callback is given with the keyword argument `cb`. For example, this will print
+"training" every 10 seconds (using [`Flux.throttle`](@ref)):
+
+    train!(loss, params, data, opt, cb = throttle(() -> println("training"), 10))
+
+The callback can call [`Flux.stop`](@ref) to interrupt the training loop.
+
+Multiple optimisers and callbacks can be passed to `opt` and `cb` as arrays.
+"""
+function train!(loss, ps, data, opt; cb = () -> ())
+  ps = Params(ps)
+  cb = runall(cb)
+  @progress for d in data
+    try
+      if d isa AbstractArray{<:Number}
+        gs = gradient(ps) do
+          loss(d)
+        end
+      else
+        gs = gradient(ps) do
+          loss(d...)
+        end
+      end
+      update!(opt, ps, gs)
+      cb()
+    catch ex
+      if ex isa StopException
+        break
+      else
+        rethrow(ex)
+      end
+    end
+  end
+end
+
+"""
+    @epochs N body
+
+Run `body` `N` times. Mainly useful for quickly doing multiple epochs of
+training in a REPL.
+
+# Examples
+```jldoctest
+julia> Flux.@epochs 2 println("hello")
+[ Info: Epoch 1
+hello
+[ Info: Epoch 2
+hello
+```
+"""
+macro epochs(n, ex)
+  :(@progress for i = 1:$(esc(n))
+      @info "Epoch $i"
+      $(esc(ex))
+    end)
+end

src/tracker/Tracker.jl

@@ -1,80 +0,0 @@
-module Tracker
-
-using Base: RefValue
-
-export TrackedArray, param, back!
-
-data(x) = x
-istracked(x) = false
-
-struct Call{F,As<:Tuple}
-  func::F
-  args::As
-end
-
-Call(f, args...) = Call{typeof(f),typeof(args)}(f, args)
-
-(c::Call)() = c.func(data.(c.args)...)
-
-struct TrackedArray{T,N,A} <: AbstractArray{T,N}
-  ref::RefValue{UInt32}
-  f::Call
-  data::A
-  grad::RefValue{A}
-end
-
-TrackedScalar{T,A} = TrackedArray{T,0,A}
-TrackedVector{T,A} = TrackedArray{T,1,A}
-TrackedMatrix{T,A} = TrackedArray{T,2,A}
-TrackedVecOrMat{T,A} = Union{TrackedVector{T,A},TrackedMatrix{T,A}}
-
-TrackedArray(c::Call, x::A, Δ::Ref{A}) where A <: AbstractArray =
-  TrackedArray{eltype(A),ndims(A),A}(Ref(UInt32(0)), c, x, Δ)
-
-TrackedArray(c::Call, x::AbstractArray) = TrackedArray(c, x, RefValue{typeof(x)}())
-
-TrackedArray(c::Call) = TrackedArray(c, c())
-
-TrackedArray(x::AbstractArray) = TrackedArray(Call(nothing), x, RefValue(zeros(x)))
-
-param(xs) = TrackedArray(AbstractFloat.(xs))
-istracked(x::TrackedArray) = true
-data(x::TrackedArray) = x.data
-grad(x::TrackedArray) = x.grad[]
-
-# Fallthrough methods
-
-for f in :[Base.size, Base.ndims].args
-  @eval @inline $f(x::TrackedArray, a...) = $f(data(x), a...)
-end
-
-Base.similar(x::TrackedArray, dims::Union{AbstractUnitRange,Integer}...) =
-  similar(data(x), dims...)
-
-Base.similar(x::TrackedArray, T::Type) = similar(data(x), T)
-
-Base.show(io::IO, ::Type{TrackedArray{T,N,A}}) where {T,N,A<:AbstractArray{T,N}} =
-  print(io, "TrackedArray{…,$A}")
-
-function Base.showarray(io::IO, X::TrackedArray, repr::Bool = true; header = true)
-  if repr
-    print(io, "param(")
-    Base.showarray(io, data(X), true)
-    print(io, ")")
-  else
-    header && print(io, "Tracked ")
-    Base.showarray(io, data(X), false, header = header)
-  end
-end
-
-include("back.jl")
-include("lib.jl")
-include("numeric.jl")
-
-import NNlib.adapt
-
-adapt(T, xs::TrackedArray) =
-  TrackedArray(xs.f, adapt(T, xs.data),
-               RefValue(adapt(T, grad(xs))))
-
-end

src/tracker/back.jl

@@ -1,43 +0,0 @@
-scan(x) = nothing
-
-scan(c::Call) = foreach(scan, c.args)
-
-function scan(x::TrackedArray)
-  ref = x.ref[] += 1
-  if ref == 1
-    scan(x.f)
-  else
-    isassigned(x.grad) || (x.grad[] = zeros(x.data))
-  end
-  return
-end
-
-back(c::Call, Δ) = back(c.func, Δ, c.args...)
-back(::Call{Void}, Δ) = nothing
-
-function back(x::TrackedArray, Δ)
-  ref = x.ref[] -= 1
-  if isassigned(x.grad)
-    x.grad[] .+= Δ
-    ref == 0 && back(x.f, x.grad[])
-  else
-    ref == 0 && back(x.f, Δ)
-  end
-  return
-end
-
-macro back(x, Δ)
-  quote
-    x = $(esc(x))
-    istracked(x) && back(x, $(esc(Δ)))
-  end
-end
-
-# Interface methods
-
-function back!(x::TrackedArray, Δ)
-  scan(x)
-  back(x, Δ)
-end
-
-back!(x::TrackedScalar) = back!(x, 1)

src/tracker/lib.jl

@@ -1,129 +0,0 @@
-import Base: *
-
-toarray(xs::AbstractArray, ys::AbstractArray) = ys
-toarray(xs::AbstractArray, y) = similar(xs, typeof(y), ()) .= y
-
-unarray(xs) = xs
-unarray(xs::AbstractArray{T,0} where T) = xs[]
-
-Base.getindex(xs::TrackedArray, i...) =
-  TrackedArray(Call(getindex, xs, i...), toarray(xs.data, xs.data[i...]))
-
-function back(::typeof(getindex), Δ, xs::TrackedArray, i...)
-  Δ′ = zeros(xs.data)
-  Δ′[i...] = unarray(Δ)
-  @back(xs, Δ′)
-end
-
-Base.:-(xs::TrackedArray) = TrackedArray(Call(-, xs))
-
-back(::typeof(-), Δ, xs::TrackedArray) = back(xs, -Δ)
-
-Base.transpose(xs::TrackedArray) = TrackedArray(Call(transpose, xs))
-Base.ctranspose(xs::TrackedArray) = TrackedArray(Call(ctranspose, xs))
-
-back(::typeof(transpose), Δ, xs) = @back(xs, trim(xs, Δ.'))
-back(::typeof(ctranspose), Δ, xs) = @back(xs, trim(xs, Δ'))
-
-Base.repmat(x::TrackedVecOrMat, a::Integer...) = TrackedArray(Call(repmat, x, a...))
-Base.repmat(x::TrackedVecOrMat, a::Int64...) = TrackedArray(Call(repmat, x, a...))
-
-Base.vcat(a::TrackedVector, b::TrackedVector)  = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::TrackedVector, b::AbstractVector) = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::AbstractVector, b::TrackedVector) = TrackedArray(Call(vcat, a, b))
-
-Base.vcat(a::TrackedVecOrMat, b::TrackedVecOrMat)  = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::TrackedVecOrMat, b::AbstractVecOrMat) = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::AbstractVecOrMat, b::TrackedVecOrMat) = TrackedArray(Call(vcat, a, b))
-
-Base.vcat(a::TrackedMatrix, b::TrackedMatrix)  = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::TrackedMatrix, b::AbstractMatrix) = TrackedArray(Call(vcat, a, b))
-Base.vcat(a::AbstractMatrix, b::TrackedMatrix) = TrackedArray(Call(vcat, a, b))
-
-function back(::typeof(vcat), Δ, xs, ys)
-  i = Base.tail(map(_ -> :, size(Δ)))
-  @back(xs, Δ[1:size(xs,1), i...])
-  @back(ys, Δ[size(xs,1)+1:end, i...])
-end
-
-# Reductions
-
-Base.sum(xs::TrackedArray, dim) = TrackedArray(Call(sum, xs, dim))
-Base.sum(xs::TrackedArray) = TrackedArray(Call(sum, xs), toarray(xs.data, sum(xs.data)))
-Base.sum(xs::TrackedScalar, dim...) = xs
-
-back(::typeof(sum), Δ, xs::TrackedArray, dim...) = back(xs, similar(xs.data) .= Δ)
-
-Base.maximum(xs::TrackedArray, args...) = maximum(xs.data, args...)
-Base.findfirst(xs::TrackedArray, args...) = findfirst(xs.data, args...)
-
-# BLAS
-
-a::TrackedMatrix * b::TrackedMatrix  = TrackedArray(Call(*, a, b))
-a::TrackedMatrix * b::AbstractMatrix = TrackedArray(Call(*, a, b))
-a::AbstractMatrix * b::TrackedMatrix = TrackedArray(Call(*, a, b))
-
-a::TrackedMatrix * b::TrackedVector  = TrackedArray(Call(*, a, b))
-a::TrackedMatrix * b::AbstractVector = TrackedArray(Call(*, a, b))
-a::AbstractMatrix * b::TrackedVector = TrackedArray(Call(*, a, b))
-
-function back(::typeof(*), Δ, a::AbstractMatrix, b::AbstractVecOrMat)
-  @back(a, A_mul_Bt(Δ, data(b)))
-  @back(b, At_mul_B(data(a), Δ))
-end
-
-# NNlib
-
-import NNlib: softmax, ∇softmax
-
-softmax(xs::TrackedArray) = TrackedArray(Call(softmax, xs))
-
-back(::typeof(softmax), Δ, xs) = @back(xs, ∇softmax(Δ, data(xs)))
-
-# Broadcasting
-
-using ForwardDiff: Dual, partials
-
-struct Broadcasted{T}
-  data::T
-end
-
-(b::Broadcasted)(xs...) = map(x -> x.value, b.data)
-
-dualify(xs, n) = xs
-dualify(xs::TrackedArray, ps) = map(x -> Dual(x, ps), data(xs))
-
-function tracked_broadcast(f, args::Vararg{Any,N}) where N
-  dargs = map((x,i) -> dualify(x, ntuple(j -> i==j, Val{N})), args, ntuple(identity, Val{N}))
-  # TrackedArray(Call(Broadcasted(broadcast(f, dargs...)), args...))
-  # Works around a 0.6 type inference issue
-  b = Broadcasted(broadcast(f, dargs...))
-  TrackedArray(Call(b, args...), b())
-end
-
-trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val{ndims(x)}))
-
-unbroadcast(x, Δ) =
-  size(x) == size(Δ) ? Δ :
-    trim(x, sum(Δ, filter(n -> size(x, n) == 1, 1:ndims(Δ))))
-
-function getpartial(Δ, x, i)
-  @inbounds p = getindex(partials(x), i)
-  return Δ * p
-end
-
-function back(b::Broadcasted, Δ, args::Vararg{Any,N}) where N
-  Δargs = ntuple(i -> getpartial.(Δ, b.data, i), Val{N})
-  foreach((x, Δ) -> @back(x, unbroadcast(x, Δ)), args, Δargs)
-end
-
-Base.Broadcast._containertype(::Type{<:TrackedArray}) = TrackedArray
-Base.Broadcast.promote_containertype(::Type{TrackedArray}, ::Type{TrackedArray}) = TrackedArray
-Base.Broadcast.promote_containertype(::Type{Array}, ::Type{TrackedArray}) = TrackedArray
-Base.Broadcast.promote_containertype(::Type{TrackedArray}, ::Type{Array}) = TrackedArray
-Base.Broadcast.promote_containertype(::Type{TrackedArray}, ct) = TrackedArray
-Base.Broadcast.promote_containertype(ct, ::Type{TrackedArray}) = TrackedArray
-Base.Broadcast.broadcast_indices(::Type{TrackedArray}, A::Ref) = ()
-Base.Broadcast.broadcast_indices(::Type{TrackedArray}, A) = indices(A)
-
-Base.Broadcast.broadcast_c(f, ::Type{TrackedArray}, A, Bs...) = tracked_broadcast(f, A, Bs...)

src/tracker/numeric.jl

@@ -1,22 +0,0 @@
-function gradient(f, xs::AbstractArray...)
-  xs = param.(xs)
-  back!(f(xs...))
-  grad.(xs)
-end
-
-function ngradient(f, xs::AbstractArray...)
-  grads = zeros.(xs)
-  for (x, Δ) in zip(xs, grads), i in 1:length(x)
-    δ = sqrt(eps())
-    tmp = x[i]
-    x[i] = tmp - δ/2
-    y1 = f(xs...)
-    x[i] = tmp + δ/2
-    y2 = f(xs...)
-    x[i] = tmp
-    Δ[i] = (y2-y1)/δ
-  end
-  return grads
-end
-
-gradcheck(f, xs...) = all(isapprox.(ngradient(f, xs...), gradient(f, xs...), rtol = 1e-6))

src/tree.jl

@@ -1,34 +0,0 @@
-children(x) = ()
-mapchildren(f, x) = x
-
-function treelike(T, fs = fieldnames(T))
-  @eval begin
-    children(x::$T) = ($([:(x.$f) for f in fs]...),)
-    mapchildren(f, x::$T) = $T(f.(children(x))...)
-  end
-end
-
-isleaf(x) = isempty(children(x))
-
-function mapleaves(f, x; cache = ObjectIdDict())
-  haskey(cache, x) && return cache[x]
-  cache[x] = isleaf(x) ? f(x) : mapchildren(x -> mapleaves(f, x, cache = cache), x)
-end
-
-export mapparams
-@deprecate mapparams(f, x) mapleaves(f, x)
-
-using DataFlow: OSet
-
-function forleaves(f, x; seen = OSet())
-  x ∈ seen && return
-  push!(seen, x)
-  isleaf(x) ? f(x) : foreach(x -> forleaves(f, x, seen = seen), children(x))
-  return
-end
-
-function params(m)
-  ps = []
-  forleaves(p -> p isa TrackedArray && push!(ps, p), m)
-  return ps
-end

src/utils.jl

@@ -1,67 +1,315 @@
 # Arrays
+nfan() = 1, 1 # fan_in, fan_out
+nfan(n) = 1, n # A vector is treated as a n×1 matrix
+nfan(n_out, n_in) = n_in, n_out # In case of Dense kernels: arranged as matrices
+nfan(dims...) = prod(dims[1:end-2]) .* (dims[end-1], dims[end]) # In case of convolution kernels
 
-initn(dims...) = randn(dims...)/100
+"""
+    glorot_uniform(dims...)
 
-flatten(xs) = reshape(xs, size(xs, 1), :)
+Return an `Array` of size `dims` containing random variables taken from a uniform
+distribution in the interval ``[-x, x]``, where `x = sqrt(24 / sum(dims)) / 2`.
 
+# Examples
+```jldoctest; setup = :(using Random; Random.seed!(0))
+julia> Flux.glorot_uniform(2, 3)
+2×3 Array{Float32,2}:
+ 0.601094  -0.57414   -0.814925
+ 0.900868   0.805994   0.057514
+```
+"""
+glorot_uniform(dims...) = (rand(Float32, dims...) .- 0.5f0) .* sqrt(24.0f0 / sum(nfan(dims...)))
+
+"""
+    glorot_normal(dims...)
+
+Return an `Array` of size `dims` containing random variables taken from a normal
+distribution with mean 0 and standard deviation `sqrt(2 / sum(dims))`.
+
+# Examples
+```jldoctest; setup = :(using Random; Random.seed!(0))
+julia> Flux.glorot_normal(3, 2)
+3×2 Array{Float32,2}:
+  0.429505  -0.0852891
+  0.523935   0.371009
+ -0.223261   0.188052
+```
+"""
+glorot_normal(dims...) = randn(Float32, dims...) .* sqrt(2.0f0 / sum(nfan(dims...)))
+
+ones(T::Type, dims...) = Base.ones(T, dims...)
+zeros(T::Type, dims...) = Base.zeros(T, dims...)
+
+ones(dims...) = Base.ones(Float32, dims...)
+zeros(dims...) = Base.zeros(Float32, dims...)
+
+"""
+    unsqueeze(xs, dim)
+
+Return `xs` reshaped into an `Array` one dimensionality higher than `xs`,
+where `dim` indicates in which dimension `xs` is extended.
+
+# Examples
+```jldoctest
+julia> xs = [[1, 2], [3, 4], [5, 6]]
+3-element Array{Array{Int64,1},1}:
+ [1, 2]
+ [3, 4]
+ [5, 6]
+
+julia> Flux.unsqueeze(xs, 1)
+1×3 Array{Array{Int64,1},2}:
+ [1, 2]  [3, 4]  [5, 6]
+
+julia> Flux.unsqueeze([1 2; 3 4], 2)
+2×1×2 Array{Int64,3}:
+[:, :, 1] =
+ 1
+ 3
+
+[:, :, 2] =
+ 2
+ 4
+```
+"""
 unsqueeze(xs, dim) = reshape(xs, (size(xs)[1:dim-1]..., 1, size(xs)[dim:end]...))
 
-stack(xs, dim) = cat(dim, unsqueeze.(xs, dim)...)
-unstack(xs, dim) = [slicedim(xs, dim, i) for i = 1:size(xs, dim)]
+"""
+    stack(xs, dim)
 
-batchindex(xs, i) = (reverse(Base.tail(reverse(indices(xs))))..., i)
+Concatenate the given `Array` of `Array`s `xs` into a single `Array` along the
+given dimension `dim`.
 
+# Examples
+```jldoctest
+julia> xs = [[1, 2], [3, 4], [5, 6]]
+3-element Array{Array{Int64,1},1}:
+ [1, 2]
+ [3, 4]
+ [5, 6]
+
+julia> Flux.stack(xs, 1)
+3×2 Array{Int64,2}:
+ 1  2
+ 3  4
+ 5  6
+
+julia> cat(xs, dims=1)
+3-element Array{Array{Int64,1},1}:
+ [1, 2]
+ [3, 4]
+ [5, 6]
+```
+"""
+stack(xs, dim) = cat(unsqueeze.(xs, dim)..., dims=dim)
+
+"""
+    unstack(xs, dim)
+
+Unroll the given `xs` into an `Array` of `Array`s along the given dimension `dim`.
+
+# Examples
+```jldoctest
+julia> Flux.unstack([1 3 5 7; 2 4 6 8], 2)
+4-element Array{Array{Int64,1},1}:
+ [1, 2]
+ [3, 4]
+ [5, 6]
+ [7, 8]
+```
+"""
+unstack(xs, dim) = [copy(selectdim(xs, dim, i)) for i in 1:size(xs, dim)]
+
+"""
+    chunk(xs, n)
+
+Split `xs` into `n` parts.
+
+# Examples
+```jldoctest
+julia> Flux.chunk(1:10, 3)
+3-element Array{UnitRange{Int64},1}:
+ 1:4
+ 5:8
+ 9:10
+
+julia> Flux.chunk(collect(1:10), 3)
+3-element Array{SubArray{Int64,1,Array{Int64,1},Tuple{UnitRange{Int64}},true},1}:
+ [1, 2, 3, 4]
+ [5, 6, 7, 8]
+ [9, 10]
+```
+"""
+chunk(xs, n) = collect(Iterators.partition(xs, ceil(Int, length(xs)/n)))
+
+batchindex(xs, i) = (reverse(Base.tail(reverse(axes(xs))))..., i)
+
+"""
+    frequencies(xs)
+
+Count the number of times that each element of `xs` appears.
+
+# Examples
+```jldoctest
+julia> Flux.frequencies(['a','b','b'])
+Dict{Char,Int64} with 2 entries:
+  'a' => 1
+  'b' => 2
+```
+"""
+function frequencies(xs)
+  fs = Dict{eltype(xs),Int}()
+  for x in xs
+    fs[x] = get(fs, x, 0) + 1
+  end
+  return fs
+end
+
+head(x::Tuple) = reverse(Base.tail(reverse(x)))
+
+squeezebatch(x) = reshape(x, head(size(x)))
+
+"""
+    batch(xs)
+
+Batch the arrays in `xs` into a single array.
+
+# Examples
+```jldoctest
+julia> Flux.batch([[1,2,3],[4,5,6]])
+3×2 Array{Int64,2}:
+ 1  4
+ 2  5
+ 3  6
+```
+"""
 function batch(xs)
-  data = similar(first(xs), size(first(xs))..., length(xs))
+  data = first(xs) isa AbstractArray ?
+    similar(first(xs), size(first(xs))..., length(xs)) :
+    Vector{eltype(xs)}(undef, length(xs))
   for (i, x) in enumerate(xs)
     data[batchindex(data, i)...] = x
   end
   return data
 end
 
+"""
+Return the given sequence padded with `p` up to a maximum length of `n`.
+
+# Examples
+```jldoctest
+julia> rpad([1, 2], 4, 0)
+4-element Array{Int64,1}:
+ 1
+ 2
+ 0
+ 0
+
+julia> rpad([1, 2, 3], 2, 0)
+3-element Array{Int64,1}:
+ 1
+ 2
+ 3
+```
+"""
 Base.rpad(v::AbstractVector, n::Integer, p) = [v; fill(p, max(n - length(v), 0))]
 
-function batchseq(xs, pad, n = maximum(length(x) for x in xs))
+"""
+    batchseq(seqs, pad)
+
+Take a list of `N` sequences, and turn them into a single sequence where each
+item is a batch of `N`. Short sequences will be padded by `pad`.
+
+# Examples
+```jldoctest
+julia> Flux.batchseq([[1, 2, 3], [4, 5]], 0)
+3-element Array{Array{Int64,1},1}:
+ [1, 4]
+ [2, 5]
+ [3, 0]
+```
+"""
+function batchseq(xs, pad = nothing, n = maximum(length(x) for x in xs))
   xs_ = [rpad(x, n, pad) for x in xs]
   [batch([xs_[j][i] for j = 1:length(xs_)]) for i = 1:n]
 end
 
-# Other
+# Flattening models to weight vectors, and back
 
-function accuracy(m, data)
-  n = 0
-  correct = 0
-  for (x, y) in data
-    x, y = tobatch.((x, y))
-    n += size(x, 1)
-    correct += sum(argmax(m(x)) .== argmax(y))
+function _restructure(m, xs)
+  i = 0
+  fmap(m) do x
+    x isa AbstractArray || return x
+    x = reshape(xs[i.+(1:length(x))], size(x))
+    i += length(x)
+    return x
   end
-  return correct/n
+end
+
+@adjoint function _restructure(m, xs)
+  _restructure(m, xs), dm -> (nothing,destructure(dm)[1])
 end
 
 """
-Returns a function that when invoked, will only be triggered at most once
-during `timeout` seconds. Normally, the throttled function will run
-as much as it can, without ever going more than once per `wait` duration;
-but if you'd like to disable the execution on the leading edge, pass
-`leading=false`. To enable execution on the trailing edge, ditto.
+    destructure(m)
+
+Flatten a model's parameters into a single weight vector.
+
+    julia> m = Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
+    Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
+
+    julia> θ, re = destructure(m);
+
+    julia> θ
+    67-element Array{Float32,1}:
+    -0.1407104
+    ...
+
+The second return value `re` allows you to reconstruct the original network after making
+modifications to the weight vector (for example, with a hypernetwork).
+
+    julia> re(θ .* 2)
+    Chain(Dense(10, 5, σ), Dense(5, 2), softmax)
+"""
+function destructure(m)
+  xs = Zygote.Buffer([])
+  fmap(m) do x
+    x isa AbstractArray && push!(xs, x)
+    return x
+  end
+  return vcat(vec.(copy(xs))...), p -> _restructure(m, p)
+end
+
+# Other
+
+"""
+    throttle(f, timeout; leading=true, trailing=false)
+
+Return a function that when invoked, will only be triggered at most once
+during `timeout` seconds.
+
+Normally, the throttled function will run as much as it can, without ever
+going more than once per `wait` duration; but if you'd like to disable the
+execution on the leading edge, pass `leading=false`. To enable execution on
+the trailing edge, pass `trailing=true`.
 """
 function throttle(f, timeout; leading=true, trailing=false)
   cooldown = true
   later = nothing
+  result = nothing
 
   function throttled(args...; kwargs...)
     yield()
 
     if cooldown
       if leading
-        f(args...; kwargs...)
+        result = f(args...; kwargs...)
       else
         later = () -> f(args...; kwargs...)
       end
 
       cooldown = false
-      @schedule try
+      @async try
         while (sleep(timeout); later != nothing)
           later()
           later = nothing
@@ -70,9 +318,24 @@ function throttle(f, timeout; leading=true, trailing=false)
         cooldown = true
       end
     elseif trailing
-      later = () -> f(args...; kwargs...)
+      later = () -> (result = f(args...; kwargs...))
     end
 
-    nothing
+    return result
   end
 end
+
+"""
+    @jit ...
+
+The `@jit` annotation can be applied to any code, and the code will be compiled
+for performance.
+
+    @jit f(x) = @jit(x) + @jit(x)
+
+Note that compilation happens regardless of the `@jit` macro, so it should only
+be used for aesthetic purposes, or by recovering Python users.
+"""
+macro jit(ex)
+  esc(ex)
+end

src/zeros.jl

@@ -0,0 +1,106 @@
+import Base: +, -, *, reshape, size
+import Base.Broadcast: broadcasted, Broadcasted, BroadcastStyle
+
+"""
+    Zeros()
+    Zeros(size...)
+    Zeros(Type, size...)
+
+Acts as a stand-in for an array of zeros that can be
+used during training which is ignored by the optimisers.
+
+Useful to turn bias off for a forward pass of a layer.
+
+## Examples
+
+```julia
+julia> Flux.Zeros(3,3)
+3×3 Flux.Zeros{Bool,2}:
+ false  false  false
+ false  false  false
+ false  false  false
+
+julia> Flux.Zeros(Float32, 3,3)
+3×3 Flux.Zeros{Float32,2}:
+ 0.0  0.0  0.0
+ 0.0  0.0  0.0
+ 0.0  0.0  0.0
+
+julia> rand(3,3) .+ Flux.Zeros()
+3×3 Array{Float64,2}:
+ 0.198739  0.490459  0.785386
+ 0.779074  0.39986   0.66383
+ 0.854981  0.447292  0.314497
+
+julia> bias_less_conv = Conv((2,2), 1=>3, bias = Flux.Zeros())
+Conv((2, 2), 1=>3)
+```
+"""
+struct Zeros{T,N} <: AbstractArray{T,N}
+  size::Tuple
+end
+
+Zeros(::Type{T}, sz...) where T = Zeros{T,length(sz)}(sz)
+Zeros(sz::Integer...) = Zeros(Bool, sz...)
+
+Base.size(xs::Zeros) = xs.size
+Base.axes(xs::Zeros) = Base.OneTo.(size(xs))
+
+Base.IndexStyle(::Type{<:Zeros}) = IndexLinear()
+
+Base.getindex(xs::Zeros{T,N}, I::Int) where {T,N} = zero(T)
+Base.getindex(xs::Zeros{T,N}, inds::Union{Base.OneTo, Base.UnitRange}) where {T,N} =
+              Zeros(T, length(inds))
+
+Base.collect(xs::Zeros{T,N}) where {T,N} = fill(zero(T), size(xs))
+
+@adjoint reshape(xs::Zeros{T}, dims...) where T =
+                reshape(xs, dims...), _ -> nothing
+
+# Define basic ops
+for f in (:+, :-)
+  @eval @inline function $f(a::Union{AbstractArray{<:Number}, Zeros}, b::Zeros)
+    @assert size(a) == size(b) throw(DimensionMismatch("dimensions must match"))
+    a
+  end
+end
+
++(a::Zeros, b::AbstractArray) = b + a
+-(a::Zeros, b::AbstractArray) = -b + a
+
+Base.copy(xs::Zeros{T,N}) where {T,N} = xs
+
+# Define broadcasting behaviour
+for op in (:+, :-)
+  @eval function broadcasted(::typeof($op), a::AbstractArray, b::Zeros)
+    bs = Broadcast.broadcast_shape(size(a), size(b))
+    size(a) == bs && return a
+    sz = similar(a, bs)
+    sz .= a
+  end
+end
+
+broadcasted(::typeof(+), a::Zeros, b::AbstractArray) = broadcasted(+, b, a)
+broadcasted(::typeof(-), a::Zeros, b::AbstractArray) = broadcasted(+, -b, a)
+
+function broadcasted(::typeof(*), a::AbstractArray, b::Zeros)
+  Zeros(Broadcast.broadcast_shape(size(a), size(b))...)
+end
+
+broadcasted(::typeof(*), a::Zeros, b::AbstractArray) = broadcasted(*, b, a)
+
+for op in (:+, :-, :*)
+  @eval broadcasted(::typeof($op), a::Zeros, b::Zeros) = Zeros(Broadcast.broadcast_shape(size(a), size(b))...)
+end
+
+# Some opportunities to avoid scalar indexing, intermediaries
+# Since it replicates a little of what we expect Base to do,
+# it should be possible to remove in the future, but for now,
+# these help with performance.
+broadcasted(::typeof(+), a::AbstractArray, b::Zeros{T,0}) where T = a
+broadcasted(::typeof(+), a::Zeros{T,0}, b::AbstractArray) where T = b
+broadcasted(::typeof(-), a::AbstractArray, b::Zeros{T,0}) where T = a
+broadcasted(::typeof(-), a::Zeros{T,0}, b::AbstractArray) where T = -b
+broadcasted(::typeof(*), a::AbstractArray, b::Zeros{T,0}) where T = zero(a)
+broadcasted(::typeof(*), a::Zeros{T,0}, b::AbstractArray) where T = zero(b)
+broadcasted(::typeof(/), a::Zeros{T,0}, b::AbstractArray) where T = zero(b)

test/cuda/cuda.jl

@@ -0,0 +1,75 @@
+using Flux, Test
+using Flux.CuArrays
+using Flux: gpu
+
+@info "Testing GPU Support"
+
+@testset "CuArrays" begin
+
+CuArrays.allowscalar(false)
+
+x = randn(5, 5)
+cx = gpu(x)
+@test cx isa CuArray
+
+@test Flux.onecold(gpu([1.0, 2.0, 3.0])) == 3
+
+x = Flux.onehotbatch([1, 2, 3], 1:3)
+cx = gpu(x)
+@test cx isa Flux.OneHotMatrix && cx.data isa CuArray
+@test (cx .+ 1) isa CuArray
+
+m = Chain(Dense(10, 5, tanh), Dense(5, 2), softmax)
+cm = gpu(m)
+
+@test all(p isa CuArray for p in params(cm))
+@test cm(gpu(rand(10, 10))) isa CuArray{Float32,2}
+
+x = [1.,2.,3.]
+cx = gpu(x)
+@test Flux.crossentropy(x,x) ≈ Flux.crossentropy(cx,cx)
+@test Flux.crossentropy(x,x, weight=1.0) ≈ Flux.crossentropy(cx,cx, weight=1.0)
+@test Flux.crossentropy(x,x, weight=[1.0;2.0;3.0]) ≈ Flux.crossentropy(cx,cx, weight=cu([1.0;2.0;3.0]))
+
+x = [-1.1491, 0.8619, 0.3127]
+y = [1, 1, 0.]
+@test Flux.binarycrossentropy.(σ.(x),y) ≈ Array(Flux.binarycrossentropy.(cu(σ.(x)),cu(y)))
+@test Flux.logitbinarycrossentropy.(x,y) ≈ Array(Flux.logitbinarycrossentropy.(cu(x),cu(y)))
+
+xs = rand(5, 5)
+ys = Flux.onehotbatch(1:5,1:5)
+@test collect(cu(xs) .+ cu(ys)) ≈ collect(xs .+ ys)
+
+c = gpu(Conv((2,2),3=>4))
+x = gpu(rand(10, 10, 3, 2))
+l = c(gpu(rand(10,10,3,2)))
+@test gradient(x -> sum(c(x)), x)[1] isa CuArray
+
+c = gpu(CrossCor((2,2),3=>4))
+x = gpu(rand(10, 10, 3, 2))
+l = c(gpu(rand(10,10,3,2)))
+@test gradient(x -> sum(c(x)), x)[1] isa CuArray
+
+end
+
+@testset "onecold gpu" begin
+  y = Flux.onehotbatch(ones(3), 1:10) |> gpu;
+  @test Flux.onecold(y) isa CuArray
+  @test y[3,:] isa CuArray
+end
+
+@testset "restructure gpu" begin
+  dudt = Dense(1,1) |> gpu
+  p,re = Flux.destructure(dudt)
+  foo(x) = sum(re(p)(x))
+  @test gradient(foo, cu(rand(1)))[1] isa CuArray
+end
+
+if CuArrays.has_cudnn()
+  @info "Testing Flux/CUDNN"
+  include("cudnn.jl")
+  include("curnn.jl")
+  include("layers.jl")
+else
+  @warn "CUDNN unavailable, not testing GPU DNN support"
+end

test/cuda/cudnn.jl

@@ -0,0 +1,44 @@
+using Flux, CuArrays, Test
+using Flux: pullback
+
+@testset "CUDNN BatchNorm" begin
+    @testset "4D Input" begin
+        x = Float64.(collect(reshape(1:12, 2, 2, 3, 1)))
+        m = BatchNorm(3)
+        cx = gpu(x)
+        cm = gpu(m)
+
+        y, back = pullback((m, x) -> m(x), m, x)
+        cy, cback = pullback((m, x) -> m(x), cm, cx)
+
+        @test cpu(cy) ≈ y
+
+        Δ = randn(size(y))
+        dm, dx = back(Δ)
+        cdm, cdx = cback(gpu(Δ))
+
+        @test dm[].γ ≈ cpu(cdm[].γ)
+        @test dm[].β ≈ cpu(cdm[].β)
+        @test dx ≈ cpu(cdx)
+    end
+
+    @testset "2D Input" begin
+        x = Float64.(collect(reshape(1:12, 3, 4)))
+        m = BatchNorm(3)
+        cx = gpu(x)
+        cm = gpu(m)
+
+        y, back = pullback((m, x) -> m(x), m, x)
+        cy, cback = pullback((m, x) -> m(x), cm, cx)
+
+        @test cpu(cy) ≈ y
+
+        Δ = randn(size(y))
+        dm, dx = back(Δ)
+        cdm, cdx = cback(gpu(Δ))
+
+        @test dm[].γ ≈ cpu(cdm[].γ)
+        @test dm[].β ≈ cpu(cdm[].β)
+        @test dx ≈ cpu(cdx)
+    end
+end

test/cuda/curnn.jl

@@ -0,0 +1,63 @@
+using Flux, CuArrays, Test
+using Flux: pullback
+
+@testset for R in [RNN, GRU, LSTM]
+  m = R(10, 5) |> gpu
+  x = gpu(rand(10))
+  (m̄,) = gradient(m -> sum(m(x)), m)
+  Flux.reset!(m)
+  θ = gradient(() -> sum(m(x)), params(m))
+  @test collect(m̄[].cell[].Wi) == collect(θ[m.cell.Wi])
+end
+
+@testset "RNN" begin
+  @testset for R in [RNN, GRU, LSTM], batch_size in (1, 5)
+    rnn = R(10, 5)
+    curnn = fmap(gpu, rnn)
+
+    Flux.reset!(rnn)
+    Flux.reset!(curnn)
+    x = batch_size == 1 ?
+      rand(10) :
+      rand(10, batch_size)
+    cux = gpu(x)
+
+    y, back = pullback((r, x) -> r(x), rnn, x)
+    cuy, cuback = pullback((r, x) -> r(x), curnn, cux)
+
+    @test y ≈ collect(cuy)
+    @test haskey(Flux.CUDA.descs, curnn.cell)
+
+    ȳ = randn(size(y))
+    m̄, x̄ = back(ȳ)
+    cum̄, cux̄ = cuback(gpu(ȳ))
+
+    m̄[].cell[].Wi
+
+    m̄[].state
+    cum̄[].state
+
+    @test x̄ ≈ collect(cux̄)
+    @test m̄[].cell[].Wi ≈ collect(cum̄[].cell[].Wi)
+    @test m̄[].cell[].Wh ≈ collect(cum̄[].cell[].Wh)
+    @test m̄[].cell[].b ≈ collect(cum̄[].cell[].b)
+    if m̄[].state isa Tuple
+      for (x, cx) in zip(m̄[].state, cum̄[].state)
+        @test x ≈ collect(cx)
+      end
+    else
+      @test m̄[].state ≈ collect(cum̄[].state)
+    end
+
+    Flux.reset!(rnn)
+    Flux.reset!(curnn)
+    ohx = batch_size == 1 ?
+      Flux.onehot(rand(1:10), 1:10) :
+      Flux.onehotbatch(rand(1:10, batch_size), 1:10)
+    cuohx = gpu(ohx)
+    y = (rnn(ohx); rnn(ohx))
+    cuy = (curnn(cuohx); curnn(cuohx))
+
+    @test y ≈ collect(cuy)
+  end
+end

test/cuda/layers.jl

@@ -0,0 +1,98 @@
+# Test layers and data/model movements on and off the GPU
+# Add tests for layers and their gradients on the GPU
+# Most of the forward passes should be fine being applied
+# to bitstype objects, but this gives higher coverage for our use-cases
+# Check that getting the gradients does not throw
+
+# generic movement tests
+@testset "Basic GPU Movement" begin
+  @test gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple
+  @test gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple
+end
+
+# TODO: These layers get into scalar indexing
+# `AlphaDropout` throws a compilation error on GPUs,
+# whereas, the rest are scalar indexing issues.
+const BROKEN_LAYERS = [DepthwiseConv,
+		       AlphaDropout,
+                       InstanceNorm,
+                       GroupNorm]
+
+function gradtest(name::String, layers::Vector, xs = nothing, args...)
+  isnothing(xs) && error("Missing input to test the layers against.")
+  @testset "$name GPU grad tests" begin
+    for layer in layers
+      @testset "$layer GPU grad test" begin
+        l = gpu(layer(args...))
+        xs = gpu(xs)
+        if any(x -> isa(l, x), BROKEN_LAYERS)
+          ps = Flux.params(l)
+          @test_broken gradient(() -> sum(l(xs)), ps) isa Flux.Zygote.Grads
+        else
+          ps = Flux.params(l)
+          @test gradient(() -> sum(l(xs)), ps) isa Flux.Zygote.Grads
+          gs = gradient(() -> sum(l(xs)), ps)
+
+          # Handle pooling layers
+          if !isempty(ps)
+            @test gs[first(ps)] isa Flux.CuArrays.CuArray
+          end
+        end
+      end
+    end
+  end
+end
+
+# Repeats from Conv, CrossCor
+
+r = rand(Float32, 28, 28, 1, 1)
+conv_layers = [Conv, ConvTranspose, CrossCor, DepthwiseConv]
+gradtest("Conv", conv_layers, r, (2,2), 1=>3)
+
+pooling_layers = [MaxPool, MeanPool]
+gradtest("Pooling", pooling_layers, r, (2,2))
+
+dropout_layers = [Dropout, AlphaDropout]
+gradtest("Dropout", dropout_layers, r, 0.5f0)
+
+norm_layers = [LayerNorm, BatchNorm]
+gradtest("Normalising", norm_layers, rand(Float32, 28,28,3,1), 1)
+
+instancenorm = [InstanceNorm]
+gradtest("InstanceNorm", instancenorm, r, 1)
+
+groupnorm = [GroupNorm]
+gradtest("GroupNorm", groupnorm, rand(Float32, 28,28,3,1), 3, 1)
+
+const stateless_layers = [Flux.mse,
+                          Flux.crossentropy,
+                          Flux.logitcrossentropy,
+                          Flux.normalise]
+
+const stateless_layers_broadcasted = [Flux.binarycrossentropy,
+                                      Flux.logitbinarycrossentropy]
+
+function stateless_gradtest(f, args...)
+  @test gradient((args...) -> sum(f(args...)), args...)[1] isa CuArray
+end
+
+function stateless_gradtest_broadcasted(f, args...)
+  @test gradient((args...) -> sum(f.(args...)), args...)[1] isa CuArray
+end
+
+@testset "Stateless GPU grad tests" begin
+  x = gpu(rand(3,3))
+  y = gpu(rand(3,3))
+
+  for layer in stateless_layers
+    if layer == Flux.normalise
+      stateless_gradtest(layer, x)
+    else
+      stateless_gradtest(layer, x, y)
+    end
+  end
+
+  for layer in stateless_layers_broadcasted
+    stateless_gradtest_broadcasted(layer, x, y)
+  end
+end

test/data.jl

@@ -0,0 +1,116 @@
+@testset "DataLoader" begin
+    X = reshape([1:10;], (2, 5))
+    Y = [1:5;]
+
+    d = DataLoader(X, batchsize=2)
+    @inferred first(d)
+    batches = collect(d)
+    @test eltype(batches) == eltype(d) == typeof(X)
+    @test length(batches) == 3
+    @test batches[1] == X[:,1:2]
+    @test batches[2] == X[:,3:4]
+    @test batches[3] == X[:,5:5]
+
+    d = DataLoader(X, batchsize=2, partial=false)
+    @inferred first(d)
+    batches = collect(d)
+    @test eltype(batches) == eltype(d) == typeof(X)
+    @test length(batches) == 2
+    @test batches[1] == X[:,1:2]
+    @test batches[2] == X[:,3:4]
+
+    d = DataLoader((X,), batchsize=2, partial=false)
+    @inferred first(d)
+    batches = collect(d)
+    @test eltype(batches) == eltype(d) == Tuple{typeof(X)}
+    @test length(batches) == 2
+    @test batches[1] == (X[:,1:2],)
+    @test batches[2] == (X[:,3:4],)
+
+    d = DataLoader((X, Y), batchsize=2)
+    @inferred first(d)
+    batches = collect(d)
+    @test eltype(batches) == eltype(d) == Tuple{typeof(X), typeof(Y)}
+    @test length(batches) == 3
+    @test length(batches[1]) == 2
+    @test length(batches[2]) == 2
+    @test length(batches[3]) == 2
+    @test batches[1][1] == X[:,1:2]
+    @test batches[1][2] == Y[1:2]
+    @test batches[2][1] == X[:,3:4]
+    @test batches[2][2] == Y[3:4]
+    @test batches[3][1] == X[:,5:5]
+    @test batches[3][2] == Y[5:5]
+
+    # test with NamedTuple
+    d = DataLoader((x=X, y=Y), batchsize=2)
+    @inferred first(d)
+    batches = collect(d)
+    @test eltype(batches) == eltype(d) == NamedTuple{(:x, :y), Tuple{typeof(X), typeof(Y)}}
+    @test length(batches) == 3
+    @test length(batches[1]) == 2
+    @test length(batches[2]) == 2
+    @test length(batches[3]) == 2
+    @test batches[1][1] == batches[1].x == X[:,1:2]
+    @test batches[1][2] == batches[1].y == Y[1:2]
+    @test batches[2][1] == batches[2].x == X[:,3:4]
+    @test batches[2][2] == batches[2].y == Y[3:4]
+    @test batches[3][1] == batches[3].x == X[:,5:5]
+    @test batches[3][2] == batches[3].y == Y[5:5]
+
+    # test interaction with `train!`
+    θ = ones(2)
+    X = zeros(2, 10)
+    loss(x) = sum((x .- θ).^2)
+    d  = DataLoader(X) 
+    Flux.train!(loss, [θ], ncycle(d, 10), Descent(0.1))
+    @test norm(θ) < 1e-4
+
+    # test interaction with `train!`
+    θ = zeros(2)
+    X = ones(2, 10)
+    Y = fill(2, 10)
+    loss(x, y) = sum((y - x'*θ).^2)
+    d  = DataLoader((X, Y)) 
+    Flux.train!(loss, [θ], ncycle(d, 10), Descent(0.1))
+    @test norm(θ .- 1) < 1e-10
+end
+
+@testset "CMUDict" begin 
+    @test cmudict()["CATASTROPHE"] == :[K,AH0,T,AE1,S,T,R,AH0,F,IY0].args
+
+    @test length(CMUDict.phones()) == 39
+
+    @test length(CMUDict.symbols()) == 84
+end
+
+@testset "MNIST" begin 
+    @test MNIST.images()[1] isa Matrix
+    @test MNIST.labels() isa Vector{Int64}
+end
+
+@testset "FashionMNIST" begin 
+    @test FashionMNIST.images()[1] isa Matrix
+    @test FashionMNIST.labels() isa Vector{Int64}
+end
+
+@testset "Sentiment" begin 
+    @test Data.Sentiment.train() isa Vector{Data.Tree{Any}}
+end
+
+@testset "Iris" begin 
+    @test Iris.features() isa Matrix
+    @test size(Iris.features()) == (4,150)
+
+    @test Iris.labels() isa Vector{String}
+    @test size(Iris.labels()) == (150,)
+end
+
+
+@testset "Housing" begin
+    @test Housing.features() isa Matrix # test broken due to SSL certifate expiration problem
+    @test size(Housing.features()) == (506, 13)
+
+    @test Housing.targets() isa Array{Float64}
+    @test size(Housing.targets()) == (506, 1)
+end

test/layers/basic.jl

@@ -0,0 +1,117 @@
+using Test, Random
+import Flux: activations
+
+@testset "basic" begin
+  @testset "helpers" begin
+    @testset "activations" begin
+      dummy_model = Chain(x->x.^2, x->x .- 3, x -> tan.(x))
+      x = randn(10)
+      @test activations(dummy_model, x)[1] == x.^2
+      @test activations(dummy_model, x)[2] == (x.^2 .- 3)
+      @test activations(dummy_model, x)[3] == tan.(x.^2 .- 3)
+
+      @test activations(Chain(), x) == ()
+      @test activations(Chain(identity, x->:foo), x)[2] == :foo # results include `Any` type
+    end
+  end
+
+  @testset "Chain" begin
+    @test_nowarn Chain(Dense(10, 5, σ), Dense(5, 2))(randn(10))
+    @test_throws DimensionMismatch Chain(Dense(10, 5, σ),Dense(2, 1))(randn(10))
+    # numeric test should be put into testset of corresponding layer
+  end
+
+  @testset "Activations" begin
+    c = Chain(Dense(3,5,relu), Dense(5,1,relu))
+    X = Float32.([1.0; 1.0; 1.0])
+    @test_nowarn gradient(()->Flux.activations(c, X)[2][1], params(c))
+  end
+
+  @testset "Dense" begin
+    @testset "constructors" begin
+      @test size(Dense(10, 100).W) == (100, 10)
+      @test Dense(rand(100,10), rand(10)).σ == identity
+
+      @test_throws MethodError Dense(10, 10.5)
+      @test_throws MethodError Dense(10, 10.5, tanh)
+    end
+
+    @test  length(Dense(10, 5)(randn(10))) == 5
+    @test_throws DimensionMismatch Dense(10, 5)(randn(1))
+    @test_throws MethodError Dense(10, 5)(1) # avoid broadcasting
+    @test_throws MethodError Dense(10, 5).(randn(10)) # avoid broadcasting
+
+    @test Dense(10, 1, identity, initW = ones, initb = zeros)(ones(10,1)) == 10*ones(1, 1)
+    @test Dense(10, 1, identity, initW = ones, initb = zeros)(ones(10,2)) == 10*ones(1, 2)
+    @test Dense(10, 2, identity, initW = ones, initb = zeros)(ones(10,1)) == 10*ones(2, 1)
+    @test Dense(10, 2, identity, initW = ones, initb = zeros)([ones(10,1) 2*ones(10,1)]) == [10 20; 10 20]
+  end
+
+  @testset "Diagonal" begin
+    @test length(Flux.Diagonal(10)(randn(10))) == 10
+    @test length(Flux.Diagonal(10)(1)) == 10
+    @test length(Flux.Diagonal(10)(randn(1))) == 10
+    @test_throws DimensionMismatch Flux.Diagonal(10)(randn(2))
+
+    @test Flux.Diagonal(2)([1 2]) == [1 2; 1 2]
+    @test Flux.Diagonal(2)([1,2]) == [1,2]
+    @test Flux.Diagonal(2)([1 2; 3 4]) == [1 2; 3 4]
+  end
+
+  @testset "Maxout" begin
+    # Note that the normal common usage of Maxout is as per the docstring
+    # These are abnormal constructors used for testing purposes
+
+    @testset "Constructor" begin
+      mo = Maxout(() -> identity, 4)
+      input = rand(40)
+      @test mo(input) == input
+    end
+
+    @testset "simple alternatives" begin
+      mo = Maxout((x -> x, x -> 2x, x -> 0.5x))
+      input = rand(40)
+      @test mo(input) == 2*input
+    end
+
+    @testset "complex alternatives" begin
+      mo = Maxout((x -> [0.5; 0.1]*x, x -> [0.2; 0.7]*x))
+      input = [3.0 2.0]
+      target = [0.5, 0.7].*input
+      @test mo(input) == target
+    end
+
+    @testset "params" begin
+      mo = Maxout(()->Dense(32, 64), 4)
+      ps = params(mo)
+      @test length(ps) == 8  #4 alts, each with weight and bias
+    end
+  end
+
+  @testset "SkipConnection" begin
+    @testset "zero sum" begin
+      input = randn(10, 10, 10, 10)
+      @test SkipConnection(x -> zeros(size(x)), (a,b) -> a + b)(input) == input
+    end
+
+    @testset "concat size" begin
+      input = randn(10, 2)
+      @test size(SkipConnection(Dense(10,10), (a,b) -> cat(a, b, dims = 2))(input)) == (10,4)
+    end
+  end
+
+  @testset "output dimensions" begin
+    m = Chain(Conv((3, 3), 3 => 16), Conv((3, 3), 16 => 32))
+    @test Flux.outdims(m, (10, 10)) == (6, 6)
+
+    m = Dense(10, 5)
+    @test Flux.outdims(m, (5, 2)) == (5,)
+    @test Flux.outdims(m, (10,)) == (5,)
+
+    m = Flux.Diagonal(10)
+    @test Flux.outdims(m, (10,)) == (10,)
+
+    m = Maxout(() -> Conv((3, 3), 3 => 16), 2)
+    @test Flux.outdims(m, (10, 10)) == (8, 8)
+  end
+end

test/layers/conv.jl

@@ -0,0 +1,218 @@
+using Flux, Test
+using Flux: maxpool, meanpool
+using Flux: gradient
+
+@testset "Pooling" begin
+  x = randn(Float32, 10, 10, 3, 2)
+  gmp = GlobalMaxPool()
+  @test size(gmp(x)) == (1, 1, 3, 2)
+  gmp = GlobalMeanPool()
+  @test size(gmp(x)) == (1, 1, 3, 2)
+  mp = MaxPool((2, 2))
+  @test mp(x) == maxpool(x, PoolDims(x, 2))
+  mp = MeanPool((2, 2))
+  @test mp(x) == meanpool(x, PoolDims(x, 2))
+end
+
+@testset "CNN" begin
+  r = zeros(Float32, 28, 28, 1, 5)
+  m = Chain(
+    Conv((2, 2), 1=>16, relu),
+    MaxPool((2,2)),
+    Conv((2, 2), 16=>8, relu),
+    MaxPool((2,2)),
+    x -> reshape(x, :, size(x, 4)),
+    Dense(288, 10), softmax)
+
+  @test size(m(r)) == (10, 5)
+
+  # Test bias switch
+  bias = Conv(ones(Float32, 2, 2, 1, 3), ones(Float32, 3))
+  ip = zeros(Float32, 28,28,1,1)
+
+  op = bias(ip)
+  @test sum(op) == prod(size(op))
+
+  bias = Conv((2,2), 1=>3, bias = Flux.Zeros())
+  op = bias(ip)
+  @test sum(op) === 0.f0
+  gs = gradient(() -> sum(bias(ip)), Flux.params(bias))
+  @test gs[bias.bias] == nothing
+
+  # Train w/o bias and make sure no convergence happens
+  # when only bias can be converged
+  bias = Conv((2, 2), 1=>3, bias = Flux.Zeros());
+  ip = zeros(Float32, 28,28,1,1)
+  op = zeros(Float32, 27,27,3,1) .+ 2.f0
+  opt = Descent()
+
+  for _ = 1:10^3
+    gs = gradient(params(bias)) do
+      Flux.mse(bias(ip), op)
+    end
+    Flux.Optimise.update!(opt, params(bias), gs)
+  end
+
+  @test Flux.mse(bias(ip), op) ≈ 4.f0
+end
+
+@testset "asymmetric padding" begin
+  r = ones(Float32, 28, 28, 1, 1)
+  m = Conv((3, 3), 1=>1, relu; pad=(0,1,1,2))
+  m.weight[:] .= 1.0
+  m.bias[:] .= 0.0
+  y_hat = m(r)[:,:,1,1]
+  @test size(y_hat) == (27, 29)
+  @test y_hat[1, 1] ≈ 6.0
+  @test y_hat[2, 2] ≈ 9.0
+  @test y_hat[end, 1] ≈ 4.0
+  @test y_hat[1, end] ≈ 3.0
+  @test y_hat[1, end-1] ≈ 6.0
+  @test y_hat[end, end] ≈ 2.0
+end
+
+@testset "Depthwise Conv" begin
+  r = zeros(Float32, 28, 28, 3, 5)
+  m1 = DepthwiseConv((2, 2), 3=>15)
+  @test size(m1(r), 3) == 15
+
+  m3 = DepthwiseConv((2, 3), 3=>9)
+  @test size(m3(r), 3) == 9
+
+  # Test that we cannot ask for non-integer multiplication factors
+  @test_throws AssertionError DepthwiseConv((2,2), 3=>10)
+end
+
+@testset "ConvTranspose" begin
+  x = zeros(Float32, 28, 28, 1, 1)
+  y = Conv((3,3), 1 => 1)(x)
+  x_hat = ConvTranspose((3, 3), 1 => 1)(y)
+  @test size(x_hat) == size(x)
+
+  m = ConvTranspose((3,3), 1=>1)
+  # Test that the gradient call does not throw: #900
+  @test gradient(()->sum(m(x)), params(m)) isa Flux.Zygote.Grads
+end
+
+@testset "CrossCor" begin
+  x = rand(Float32, 28, 28, 1, 1)
+  w = rand(2,2,1,1)
+  y = CrossCor(w, [0.0])
+
+  @test isapprox(sum(w .* x[1:2, 1:2, :, :]), y(x)[1, 1, 1, 1], rtol=1e-7)
+
+  r = zeros(Float32, 28, 28, 1, 5)
+  m = Chain(
+    CrossCor((2, 2), 1=>16, relu),
+    MaxPool((2,2)),
+    CrossCor((2, 2), 16=>8, relu),
+    MaxPool((2,2)),
+    x -> reshape(x, :, size(x, 4)),
+    Dense(288, 10), softmax)
+
+  @test size(m(r)) == (10, 5)
+  @test y(x) != Conv(w, [0.0])(x)
+  @test CrossCor(w[end:-1:1, end:-1:1, :, :], [0.0])(x) == Conv(w, [0.0])(x)
+end
+
+@testset "Conv with non quadratic window #700" begin
+  data = zeros(Float32, 7,7,1,1)
+  data[4,4,1,1] = 1
+
+  l = Conv((3,3), 1=>1)
+  expected = zeros(eltype(l.weight),5,5,1,1)
+  expected[2:end-1,2:end-1,1,1] = l.weight
+  @test expected ≈ l(data)
+
+  l = Conv((3,1), 1=>1)
+  expected = zeros(eltype(l.weight),5,7,1,1)
+  expected[2:end-1,4,1,1] = l.weight
+  @test expected ≈ l(data)
+
+  l = Conv((1,3), 1=>1)
+  expected = zeros(eltype(l.weight),7,5,1,1)
+  expected[4,2:end-1,1,1] = l.weight
+  @test expected ≈ l(data)
+
+  @test begin
+    # we test that the next expression does not throw
+    randn(Float32, 10,10,1,1) |> Conv((6,1), 1=>1, Flux.σ)
+    true
+  end
+end
+
+@testset "conv output dimensions" begin
+  m = Conv((3, 3), 3 => 16)
+  @test Flux.outdims(m, (10, 10)) == (8, 8)
+  m = Conv((3, 3), 3 => 16; stride = 2)
+  @test Flux.outdims(m, (5, 5)) == (2, 2)
+  m = Conv((3, 3), 3 => 16; stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+  m = Conv((3, 3), 3 => 16; stride = 2, pad = 3, dilation = 2)
+  @test Flux.outdims(m, (5, 5)) == (4, 4)
+
+  m = ConvTranspose((3, 3), 3 => 16)
+  @test Flux.outdims(m, (8, 8)) == (10, 10)
+  m = ConvTranspose((3, 3), 3 => 16; stride = 2)
+  @test Flux.outdims(m, (2, 2)) == (5, 5)
+  m = ConvTranspose((3, 3), 3 => 16; stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+  m = ConvTranspose((3, 3), 3 => 16; stride = 2, pad = 3, dilation = 2)
+  @test Flux.outdims(m, (4, 4)) == (5, 5)
+
+  m = DepthwiseConv((3, 3), 3 => 6)
+  @test Flux.outdims(m, (10, 10)) == (8, 8)
+  m = DepthwiseConv((3, 3), 3 => 6; stride = 2)
+  @test Flux.outdims(m, (5, 5)) == (2, 2)
+  m = DepthwiseConv((3, 3), 3 => 6; stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+  m = DepthwiseConv((3, 3), 3 => 6; stride = 2, pad = 3, dilation = 2)
+  @test Flux.outdims(m, (5, 5)) == (4, 4)
+
+  m = CrossCor((3, 3), 3 => 16)
+  @test Flux.outdims(m, (10, 10)) == (8, 8)
+  m = CrossCor((3, 3), 3 => 16; stride = 2)
+  @test Flux.outdims(m, (5, 5)) == (2, 2)
+  m = CrossCor((3, 3), 3 => 16; stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+  m = CrossCor((3, 3), 3 => 16; stride = 2, pad = 3, dilation = 2)
+  @test Flux.outdims(m, (5, 5)) == (4, 4)
+
+  m = MaxPool((2, 2))
+  @test Flux.outdims(m, (10, 10)) == (5, 5)
+  m = MaxPool((2, 2); stride = 1)
+  @test Flux.outdims(m, (5, 5)) == (4, 4)
+  m = MaxPool((2, 2); stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+
+  m = MeanPool((2, 2))
+  @test Flux.outdims(m, (10, 10)) == (5, 5)
+  m = MeanPool((2, 2); stride = 1)
+  @test Flux.outdims(m, (5, 5)) == (4, 4)
+  m = MeanPool((2, 2); stride = 2, pad = 3)
+  @test Flux.outdims(m, (5, 5)) == (5, 5)
+end
+
+@testset "$ltype SamePad kernelsize $k" for ltype in (Conv, ConvTranspose, DepthwiseConv, CrossCor), k in ( (1,), (2,), (3,), (4,5), (6,7,8))
+  data = ones(Float32, (k .+ 3)..., 1,1)
+  l = ltype(k, 1=>1, pad=SamePad())
+  @test size(l(data)) == size(data)
+
+  l = ltype(k, 1=>1, pad=SamePad(), dilation = k .÷ 2)
+  @test size(l(data)) == size(data)
+
+  stride = 3
+  l = ltype(k, 1=>1, pad=SamePad(), stride = stride)
+  if ltype == ConvTranspose
+    @test size(l(data))[1:end-2] == stride .* size(data)[1:end-2] .- stride .+ 1
+  else
+    @test size(l(data))[1:end-2] == ceil.(Int, size(data)[1:end-2] ./ stride)
+  end
+end
+
+@testset "$ltype SamePad windowsize $k" for ltype in (MeanPool, MaxPool), k in ( (1,), (2,), (3,), (4,5), (6,7,8))
+  data = ones(Float32, (k .+ 3)..., 1,1)
+
+  l = ltype(k, pad=SamePad())
+  @test size(l(data))[1:end-2] == ceil.(Int, size(data)[1:end-2] ./ k)
+end

test/layers/normalisation.jl

@@ -0,0 +1,296 @@
+using Flux, Test, Statistics
+using Zygote: pullback
+
+evalwgrad(f, x...) = pullback(f, x...)[1]
+
+@testset "Dropout" begin
+  x = [1.,2.,3.]
+  @test x == Dropout(0.1)(x)
+  @test x == evalwgrad(Dropout(0), x)
+  @test zero(x) == evalwgrad(Dropout(1), x)
+
+  x = rand(100)
+  m = Dropout(0.9)
+  y = evalwgrad(m, x)
+  @test count(a->a==0, y) > 50
+  testmode!(m, true)
+  y = evalwgrad(m, x) # should override istraining
+  @test count(a->a==0, y) == 0
+  testmode!(m, false)
+  y = evalwgrad(m, x)
+  @test count(a->a==0, y) > 50
+
+  x = rand(Float32, 100)
+  m = Chain(Dense(100,100),
+            Dropout(0.9))
+  y = evalwgrad(m, x)
+  @test count(a->a == 0, y) > 50
+  testmode!(m, true)
+  y = evalwgrad(m, x) # should override istraining
+  @test count(a->a == 0, y) == 0
+
+  x = rand(100, 50)
+  m = Dropout(0.5, dims = 2)
+  y = m(x)
+  c = map(i->count(a->a==0, @view y[i, :]), 1:100)
+  @test minimum(c) == maximum(c)
+  m = Dropout(0.5, dims = 1)
+  y = m(x)
+  c = map(i->count(a->a==0, @view y[:, i]), 1:50)
+  @test minimum(c) == maximum(c)
+end
+
+@testset "BatchNorm" begin
+  let m = BatchNorm(2), x = [1.0 3.0 5.0;
+                             2.0 4.0 6.0]
+
+    @test length(params(m)) == 2
+
+    @test m.β == [0, 0]  # initβ(2)
+    @test m.γ == [1, 1]  # initγ(2)
+    # initial m.σ is 1
+    # initial m.μ is 0
+
+    y = evalwgrad(m, x)
+    @test isapprox(y, [-1.22474 0 1.22474; -1.22474 0 1.22474], atol = 1.0e-5)
+    # julia> x
+    #  2×3 Array{Float64,2}:
+    #  1.0  3.0  5.0
+    #  2.0  4.0  6.0
+    #
+    # μ of batch will be
+    #  (1. + 3. + 5.) / 3 = 3
+    #  (2. + 4. + 6.) / 3 = 4
+    #
+    # ∴ update rule with momentum:
+    #  .1 * 3 + 0 = .3
+    #  .1 * 4 + 0 = .4
+    @test m.μ ≈ reshape([0.3, 0.4], 2, 1)
+
+    # julia> .1 .* var(x, dims = 2, corrected=false) .* (3 / 2).+ .9 .* [1., 1.]
+    # 2×1 Array{Float64,2}:
+    #  1.3
+    #  1.3
+    @test m.σ² ≈ .1 .* var(x, dims = 2, corrected=false) .* (3 / 2).+ .9 .* [1., 1.]
+
+    x′ = m(x)
+    @test isapprox(x′[1], (1 .- 0.3) / sqrt(1.3), atol = 1.0e-5)
+  end
+
+  # with activation function
+  let m = BatchNorm(2, sigmoid), x = [1.0 3.0 5.0;
+                                      2.0 4.0 6.0]
+    y = m(x)
+    @test isapprox(y, sigmoid.((x .- m.μ) ./ sqrt.(m.σ² .+ m.ϵ)), atol = 1.0e-7)
+  end
+
+  let m = trainmode!(BatchNorm(2)), x = reshape(Float32.(1:6), 3, 2, 1)
+    y = reshape(permutedims(x, [2, 1, 3]), 2, :)
+    y = permutedims(reshape(m(y), 2, 3, 1), [2, 1, 3])
+    @test m(x) == y
+  end
+
+  let m = trainmode!(BatchNorm(2)), x = reshape(Float32.(1:12), 2, 3, 2, 1)
+    y = reshape(permutedims(x, [3, 1, 2, 4]), 2, :)
+    y = permutedims(reshape(m(y), 2, 2, 3, 1), [2, 3, 1, 4])
+    @test m(x) == y
+  end
+
+  let m = trainmode!(BatchNorm(2)), x = reshape(Float32.(1:24), 2, 2, 3, 2, 1)
+    y = reshape(permutedims(x, [4, 1, 2, 3, 5]), 2, :)
+    y = permutedims(reshape(m(y), 2, 2, 2, 3, 1), [2, 3, 4, 1, 5])
+    @test m(x) == y
+  end
+
+  let m = BatchNorm(32), x = randn(Float32, 416, 416, 32, 1);
+    m(x)
+    @test (@allocated m(x)) <  100_000_000
+  end
+end
+
+@testset "InstanceNorm" begin
+  # helper functions
+  expand_inst = (x, as) -> reshape(repeat(x, outer=[1, as[length(as)]]), as...)
+  # begin tests
+  let m = InstanceNorm(2), sizes = (3, 2, 2),
+        x = reshape(collect(1:prod(sizes)), sizes)
+
+      @test length(params(m)) == 2
+      x = Float64.(x)
+      @test m.β == [0, 0]  # initβ(2)
+      @test m.γ == [1, 1]  # initγ(2)
+      y = evalwgrad(m, x)
+
+      #julia> x
+      #[:, :, 1] =
+      # 1.0  4.0
+      # 2.0  5.0
+      # 3.0  6.0
+      #
+      #[:, :, 2] =
+      # 7.0  10.0
+      # 8.0  11.0
+      # 9.0  12.0
+      #
+      # μ will be
+      # (1. + 2. + 3.) / 3 = 2.
+      # (4. + 5. + 6.) / 3 = 5.
+      #
+      # (7. + 8. + 9.) / 3 = 8.
+      # (10. + 11. + 12.) / 3 = 11.
+      #
+      # ∴ update rule with momentum:
+      # (1. - .1) * 0 + .1 * (2. + 8.) / 2 = .5
+      # (1. - .1) * 0 + .1 * (5. + 11.) / 2 = .8
+      @test m.μ ≈ [0.5, 0.8]
+      # momentum * var * num_items / (num_items - 1) + (1 - momentum) * sigma_sq
+      # julia> reshape(mean(.1 .* var(x, dims = 1, corrected=false) .* (3 / 2), dims=3), :) .+ .9 .* 1.
+      # 2-element Array{Float64,1}:
+      #  1.
+      #  1.
+      @test m.σ² ≈ reshape(mean(.1 .* var(x, dims = 1, corrected=false) .* (3 / 2), dims=3), :) .+ .9 .* 1.
+
+      x′ = m(x)
+      @test isapprox(x′[1], (1 - 0.5) / sqrt(1. + 1f-5), atol = 1.0e-5)
+  end
+  # with activation function
+  let m = InstanceNorm(2, sigmoid), sizes = (3, 2, 2),
+      x = reshape(collect(1:prod(sizes)), sizes)
+    x = Float64.(x)
+    affine_shape = collect(sizes)
+    affine_shape[1] = 1
+
+    y = m(x)
+    @test isapprox(y, sigmoid.((x .- expand_inst(m.μ, affine_shape)) ./ sqrt.(expand_inst(m.σ², affine_shape) .+ m.ϵ)), atol = 1.0e-7)
+  end
+
+  let m = trainmode!(InstanceNorm(2)), sizes = (2, 4, 1, 2, 3),
+      x = Float32.(reshape(collect(1:prod(sizes)), sizes))
+    y = reshape(permutedims(x, [3, 1, 2, 4, 5]), :, 2, 3)
+    y = reshape(m(y), sizes...)
+    @test m(x) == y
+  end
+
+  # check that μ, σ², and the output are the correct size for higher rank tensors
+  let m = InstanceNorm(2), sizes = (5, 5, 3, 4, 2, 6),
+      x = reshape(Float32.(collect(1:prod(sizes))), sizes)
+    y = evalwgrad(m, x)
+    @test size(m.μ) == (sizes[end - 1], )
+    @test size(m.σ²) == (sizes[end - 1], )
+    @test size(y) == sizes
+  end
+
+  # show that instance norm is equal to batch norm when channel and batch dims are squashed
+  let m_inorm = trainmode!(InstanceNorm(2)), m_bnorm = trainmode!(BatchNorm(12)), sizes = (5, 5, 3, 4, 2, 6),
+      x = reshape(Float32.(collect(1:prod(sizes))), sizes)
+    @test m_inorm(x) == reshape(m_bnorm(reshape(x, (sizes[1:end - 2]..., :, 1))), sizes)
+  end
+
+  let m = InstanceNorm(32), x = randn(Float32, 416, 416, 32, 1);
+    m(x)
+    @test (@allocated m(x)) <  100_000_000
+  end
+
+end
+
+if VERSION >= v"1.1"
+@testset "GroupNorm" begin
+  # begin tests
+  squeeze(x) = dropdims(x, dims = tuple(findall(size(x) .== 1)...)) # To remove all singular dimensions
+
+  let m = GroupNorm(4,2), sizes = (3,4,2),
+        x = reshape(collect(1:prod(sizes)), sizes)
+
+      @test length(params(m)) == 2
+      x = Float64.(x)
+      @test m.β == [0, 0, 0, 0]  # initβ(32)
+      @test m.γ == [1, 1, 1, 1]  # initγ(32)
+
+      y = evalwgrad(m, x)
+
+      #julia> x
+      #[:, :, 1]  =
+      # 1.0  4.0  7.0  10.0
+      # 2.0  5.0  8.0  11.0
+      # 3.0  6.0  9.0  12.0
+      #
+      #[:, :, 2] =
+      # 13.0  16.0  19.0  22.0
+      # 14.0  17.0  20.0  23.0
+      # 15.0  18.0  21.0  24.0
+      #
+      # μ will be
+      # (1. + 2. + 3. + 4. + 5. + 6.) / 6 = 3.5
+      # (7. + 8. + 9. + 10. + 11. + 12.) / 6 = 9.5
+      #
+      # (13. + 14. + 15. + 16. + 17. + 18.) / 6 = 15.5
+      # (19. + 20. + 21. + 22. + 23. + 24.) / 6 = 21.5
+      #
+      # μ =
+      # 3.5   15.5
+      # 9.5   21.5
+      #
+      # ∴ update rule with momentum:
+      # (1. - .1) * 0 + .1 * (3.5 + 15.5) / 2 = 0.95
+      # (1. - .1) * 0 + .1 * (9.5 + 21.5) / 2 = 1.55
+      @test m.μ ≈ [0.95, 1.55]
+
+      # julia> mean(var(reshape(x,3,2,2,2),dims=(1,2)).* .1,dims=2) .+ .9*1.
+      # 2-element Array{Float64,1}:
+      #  1.25
+      #  1.25
+      @test m.σ² ≈ mean(squeeze(var(reshape(x,3,2,2,2),dims=(1,2))).*.1,dims=2) .+ .9*1.
+
+      x′ = m(x)
+      @test isapprox(x′[1], (1 - 0.95) / sqrt(1.25 + 1f-5), atol = 1.0e-5)
+  end
+  # with activation function
+  let m = GroupNorm(4,2, sigmoid), sizes = (3, 4, 2),
+      x = reshape(collect(1:prod(sizes)), sizes)
+    x = Float64.(x)
+    μ_affine_shape = ones(Int,length(sizes) + 1)
+    μ_affine_shape[end-1] = 2 # Number of groups
+
+    affine_shape = ones(Int,length(sizes) + 1)
+    affine_shape[end-2] = 2 # Channels per group
+    affine_shape[end-1] = 2 # Number of groups
+    affine_shape[1] = sizes[1]
+    affine_shape[end] = sizes[end]
+
+    og_shape = size(x)
+
+    y = m(x)
+    x_ = reshape(x,affine_shape...)
+    out = reshape(sigmoid.((x_ .- reshape(m.μ,μ_affine_shape...)) ./ sqrt.(reshape(m.σ²,μ_affine_shape...) .+ m.ϵ)),og_shape)
+    @test isapprox(y, out, atol = 1.0e-7)
+  end
+
+  let m = trainmode!(GroupNorm(2,2)), sizes = (2, 4, 1, 2, 3),
+      x = Float32.(reshape(collect(1:prod(sizes)), sizes))
+    y = reshape(permutedims(x, [3, 1, 2, 4, 5]), :, 2, 3)
+    y = reshape(m(y), sizes...)
+    @test m(x) == y
+  end
+
+  # check that μ, σ², and the output are the correct size for higher rank tensors
+  let m = GroupNorm(4,2), sizes = (5, 5, 3, 4, 4, 6),
+      x = Float32.(reshape(collect(1:prod(sizes)), sizes))
+    y = evalwgrad(m, x)
+    @test size(m.μ) == (m.G,1)
+    @test size(m.σ²) == (m.G,1)
+    @test size(y) == sizes
+  end
+
+  # show that group norm is the same as instance norm when the group size is the same as the number of channels
+  let IN = trainmode!(InstanceNorm(4)), GN = trainmode!(GroupNorm(4,4)), sizes = (2,2,3,4,5),
+      x = Float32.(reshape(collect(1:prod(sizes)), sizes))
+    @test IN(x) ≈ GN(x)
+  end
+
+  # show that group norm is the same as batch norm for a group of size 1 and batch of size 1
+  let BN = trainmode!(BatchNorm(4)), GN = trainmode!(GroupNorm(4,4)), sizes = (2,2,3,4,1),
+      x = Float32.(reshape(collect(1:prod(sizes)), sizes))
+    @test BN(x) ≈ GN(x)
+  end
+end
+end

test/layers/stateless.jl

@@ -0,0 +1,144 @@
+using Test
+using Flux: onehotbatch, mse, crossentropy, logitcrossentropy,
+            σ, binarycrossentropy, logitbinarycrossentropy, flatten,
+            xlogx, xlogy
+
+const ϵ = 1e-7
+
+@testset "xlogx & xlogy" begin
+  @test iszero(xlogx(0))
+  @test isnan(xlogx(NaN))
+  @test xlogx(2) ≈ 2.0 * log(2.0)
+  @inferred xlogx(2)
+  @inferred xlogx(0)
+
+  @test iszero(xlogy(0, 1))
+  @test isnan(xlogy(NaN, 1))
+  @test isnan(xlogy(1, NaN))
+  @test isnan(xlogy(NaN, NaN))
+  @test xlogy(2, 3) ≈ 2.0 * log(3.0)
+  @inferred xlogy(2, 3)
+  @inferred xlogy(0, 1)
+end
+
+@testset "losses" begin
+  # First, regression-style y's
+  y = [1, 1, 0, 0]
+  ŷ = [.9, .1, .1, .9]
+
+  @testset "mse" begin
+    @test mse(ŷ, y) ≈ (.1^2 + .9^2)/2
+  end
+
+  @testset "mae" begin
+    @test Flux.mae(ŷ, y) ≈ 1/2
+  end
+
+  @testset "huber_loss" begin
+    @test Flux.huber_loss(ŷ, y) ≈ 0.20500000000000002
+  end
+
+  y = [123.0,456.0,789.0]
+  ŷ = [345.0,332.0,789.0]
+  @testset "msle" begin
+    @test Flux.msle(ŷ, y) ≈ 0.38813985859136585
+  end
+
+  # Now onehot y's
+  y = onehotbatch([1, 1, 0, 0], 0:1)
+  ŷ = [.1 .9; .9 .1; .9 .1; .1 .9]'
+  v = log(.1 / .9)
+  logŷ = [v 0.0; 0.0 v; 0.0 v; v 0.0]'
+  lossvalue = 1.203972804325936
+
+  @testset "crossentropy" begin
+    @test crossentropy([0.1,0.0,0.9], [0.1,0.0,0.9]) ≈ crossentropy([0.1,0.9], [0.1,0.9])
+    @test crossentropy(ŷ, y) ≈ lossvalue
+  end
+
+  @testset "logitcrossentropy" begin
+    @test logitcrossentropy(logŷ, y) ≈ lossvalue
+  end
+
+  @testset "weighted_crossentropy" begin
+    @test crossentropy(ŷ, y, weight = ones(2)) ≈ lossvalue
+    @test crossentropy(ŷ, y, weight = [.5, .5]) ≈ lossvalue/2
+    @test crossentropy(ŷ, y, weight = [2, .5]) ≈ 1.5049660054074199
+  end
+
+  @testset "weighted_logitcrossentropy" begin
+    @test logitcrossentropy(logŷ, y, weight = ones(2)) ≈ lossvalue
+    @test logitcrossentropy(logŷ, y, weight = [.5, .5]) ≈ lossvalue/2
+    @test logitcrossentropy(logŷ, y, weight = [2, .5]) ≈ 1.5049660054074199
+  end
+
+  logŷ, y = randn(3), rand(3)
+  @testset "binarycrossentropy" begin
+    @test binarycrossentropy.(σ.(logŷ), y; ϵ=0) ≈ -y.*log.(σ.(logŷ)) - (1 .- y).*log.(1 .- σ.(logŷ))
+    @test binarycrossentropy.(σ.(logŷ), y) ≈ -y.*log.(σ.(logŷ) .+ eps.(σ.(logŷ))) - (1 .- y).*log.(1 .- σ.(logŷ) .+ eps.(σ.(logŷ)))
+  end
+
+  @testset "logitbinarycrossentropy" begin
+    @test logitbinarycrossentropy.(logŷ, y) ≈ binarycrossentropy.(σ.(logŷ), y; ϵ=0)
+  end
+
+  y = [1 2 3]
+  ŷ = [4.0 5.0 6.0]
+  @testset "kldivergence" begin
+    @test Flux.kldivergence([0.1,0.0,0.9], [0.1,0.0,0.9]) ≈ Flux.kldivergence([0.1,0.9], [0.1,0.9])
+    @test Flux.kldivergence(ŷ, y) ≈ -1.7661057888493457
+    @test Flux.kldivergence(y, y) ≈ 0
+  end
+
+  y = [1 2 3 4]
+  ŷ = [5.0 6.0 7.0 8.0]
+  @testset "hinge" begin
+    @test Flux.hinge(ŷ, y) ≈ 0
+    @test Flux.hinge(y, 0.5 .* y) ≈ 0.125
+  end
+
+  @testset "squared_hinge" begin
+    @test Flux.squared_hinge(ŷ, y) ≈ 0
+    @test Flux.squared_hinge(y, 0.5 .* y) ≈ 0.0625
+  end
+
+  y = [0.1 0.2 0.3]
+  ŷ = [0.4 0.5 0.6]
+  @testset "poisson" begin
+    @test Flux.poisson(ŷ, y) ≈ 0.6278353988097339
+    @test Flux.poisson(y, y) ≈ 0.5044459776946685
+  end
+
+  y = [1.0 0.5 0.3 2.4]
+  ŷ = [0 1.4 0.5 1.2]
+  @testset "dice_coeff_loss" begin
+    @test Flux.dice_coeff_loss(ŷ, y) ≈ 0.2799999999999999
+    @test Flux.dice_coeff_loss(y, y) ≈ 0.0
+  end
+
+  @testset "tversky_loss" begin
+    @test Flux.tversky_loss(ŷ, y) ≈ -0.06772009029345383
+    @test Flux.tversky_loss(ŷ, y, β = 0.8) ≈ -0.09490740740740744
+    @test Flux.tversky_loss(y, y) ≈ -0.5576923076923075
+  end
+
+  @testset "no spurious promotions" begin
+    for T in (Float32, Float64)
+      y = rand(T, 2)
+      ŷ = rand(T, 2)
+      for f in (mse, crossentropy, logitcrossentropy, Flux.kldivergence, Flux.hinge, Flux.poisson,
+              Flux.mae, Flux.huber_loss, Flux.msle, Flux.squared_hinge, Flux.dice_coeff_loss, Flux.tversky_loss)
+        fwd, back = Flux.pullback(f, ŷ, y)
+        @test fwd isa T
+        @test eltype(back(one(T))[1]) == T
+      end
+    end
+  end
+end
+
+@testset "helpers" begin
+  @testset "flatten" begin
+    x = randn(Float32, 10, 10, 3, 2)
+    @test size(flatten(x)) == (300, 2)
+  end
+end

test/onehot.jl

@@ -0,0 +1,19 @@
+using Flux:onecold
+using Test
+
+@testset "onecold" begin
+  a = [1, 2, 5, 3.]
+  A = [1 20 5; 2 7 6; 3 9 10; 2 1 14]
+  labels = ['A', 'B', 'C', 'D']
+
+  @test onecold(a) == 3
+  @test onecold(A) == [3, 1, 4]
+  @test onecold(a, labels) == 'C'
+  @test onecold(A, labels) == ['C', 'A', 'D']
+end
+
+@testset "onehotbatch indexing" begin
+  y = Flux.onehotbatch(ones(3), 1:10)
+  @test y[:,1] isa Flux.OneHotVector
+  @test y[:,:] isa Flux.OneHotMatrix
+end

test/optimise.jl

@@ -0,0 +1,113 @@
+using Flux.Optimise
+using Flux.Optimise: runall
+using Flux: Params, gradient
+using Test
+
+@testset "Optimise" begin
+  w = randn(10, 10)
+  @testset for opt in [ADAMW(), ADAGrad(0.1), AdaMax(), ADADelta(0.9), AMSGrad(),
+                       NADAM(), RADAM(), Descent(0.1), ADAM(), Nesterov(), RMSProp(),
+                       Momentum()]
+    w′ = randn(10, 10)
+    loss(x) = Flux.mse(w*x, w′*x)
+    for t = 1: 10^5
+      θ = Params([w′])
+      x = rand(10)
+      θ̄ = gradient(() -> loss(x), θ)
+      Optimise.update!(opt, θ, θ̄)
+    end
+    @test loss(rand(10, 10)) < 0.01
+  end
+end
+
+@testset "Optimiser" begin
+  w = randn(10, 10)
+  @testset for Opt in [InvDecay, WeightDecay, ExpDecay]
+    w′ = randn(10, 10)
+    loss(x) = Flux.mse(w*x, w′*x)
+    opt = Optimiser(Opt(), ADAM(0.001))
+    for t = 1:10^5
+      θ = Params([w′])
+      x = rand(10)
+      θ̄ = gradient(() -> loss(x), θ)
+      Optimise.update!(opt, θ, θ̄)
+    end
+    @test loss(rand(10, 10)) < 0.01
+  end
+end
+
+@testset "Training Loop" begin
+  i = 0
+  l = 1
+
+  Flux.train!(() -> (sleep(0.1); i += 1; l),
+              (),
+              Iterators.repeated((), 100),
+              Descent(),
+              cb = Flux.throttle(() -> (i > 3 && Flux.stop()), 1))
+
+  @test 3 < i < 50
+
+  # Test multiple callbacks
+  x = 0
+  fs = [() -> (), () -> x = 1]
+  cbs = runall(fs)
+  cbs()
+  @test x == 1
+end
+
+@testset "ExpDecay" begin
+
+  @testset "Sanity Check" begin
+    o = ExpDecay(0.2, 0.5, 1, 1e-3)
+    p = [0.0]
+    steps = 1:8
+    eta_expected = @. max(o.eta * 0.5 ^ steps, o.clip)
+    eta_actual = [Optimise.apply!(o, p, [1.0])[1] for _ in steps]
+    @test eta_actual == eta_expected
+  end
+
+  w = randn(10, 10)
+  o = ExpDecay(0.1, 0.1, 1000, 1e-4)
+  w1 = randn(10,10)
+  loss(x) = Flux.mse(w*x, w1*x)
+  flag = 1
+  decay_steps = []
+  for t = 1:10^5
+    prev_eta = o.eta
+    θ = Params([w1])
+    x = rand(10)
+    θ̄ = gradient(() -> loss(x), θ)
+    prev_grad = collect(θ̄[w1])
+    delta = Optimise.apply!(o, w1, θ̄[w1])
+    w1 .-= delta
+    new_eta = o.eta
+    if new_eta != prev_eta
+      push!(decay_steps, t)
+    end
+    array = fill(o.eta, size(prev_grad))
+    if array .* prev_grad != delta
+      flag = 0
+    end
+  end
+  @test flag == 1
+  # Test to check if decay happens at decay steps. Eta reaches clip value (1e-4) after 4000 steps (decay by 0.1 every 1000 steps starting at 0.1).
+  ground_truth = []
+  for i in 1:4
+    push!(ground_truth, 1000*i)  # Expected decay steps for this example.
+  end
+  @test decay_steps == ground_truth
+  @test o.eta == o.clip
+end
+
+@testset "Clipping" begin
+    w = randn(10, 10)
+    loss(x) = sum(w * x)
+    θ = Params([w])
+    x = 1000 * randn(10)
+    w̄ = gradient(() -> loss(x), θ)[w]
+    w̄_value = Optimise.apply!(ClipValue(1.0), w, copy(w̄))
+    @test all(w̄_value .<= 1)
+    w̄_norm = Optimise.apply!(ClipNorm(1.0), w, copy(w̄))
+    @test norm(w̄_norm) <= 1
+end

test/runtests.jl

@@ -1,8 +1,46 @@
-using Flux, Base.Test
+using Flux 
+using Flux.Data
+using Test 
+using Random, Statistics, LinearAlgebra
+using IterTools: ncycle
 
-@testset "Flux" begin
-
-include("utils.jl")
-include("tracker.jl")
+Random.seed!(0)
 
+@testset "Utils" begin
+  include("utils.jl")
+end
+
+@testset "Onehot" begin
+  include("onehot.jl")
+end
+
+@testset "Optimise" begin
+  include("optimise.jl")
+end
+
+@testset "Data" begin
+  include("data.jl")
+end
+
+@testset "Layers" begin
+  include("layers/basic.jl")
+  include("layers/normalisation.jl")
+  include("layers/stateless.jl")
+  include("layers/conv.jl")
+end
+
+@testset "CUDA" begin
+  if Flux.use_cuda[]
+    include("cuda/cuda.jl")
+  else
+    @warn "CUDA unavailable, not testing GPU support"
+  end
+end
+
+@static if VERSION >= v"1.4"
+  using Documenter
+  @testset "Docs" begin
+    DocMeta.setdocmeta!(Flux, :DocTestSetup, :(using Flux); recursive=true)
+    doctest(Flux)
+  end
 end

test/tracker.jl

@@ -1,30 +0,0 @@
-using Flux.Tracker, Base.Test, NNlib
-using Flux.Tracker: gradcheck
-
-gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(f(xs...)), xs...)
-gradtest(f, dims...) = gradtest(f, rand.(dims)...)
-
-@testset "Tracker" begin
-
-@test gradtest((x, W, b) -> σ.(W*x .+ b), 5, (2,5), 2)
-@test gradtest((x, W, b) -> σ.(W*x .+ b), (5,3), (2,5), 2)
-
-@test gradtest(x -> sin.(sum(x, (2, 3))), (3,4,5))
-
-@test gradtest(x -> NNlib.softmax(x).*(1:3), 3)
-@test gradtest(x -> NNlib.softmax(x).*(1:3), (3,5))
-
-@test gradtest(Flux.mse, rand(5,5), rand(5, 5))
-@test gradtest(Flux.crossentropy, rand(5,5), rand(5, 5))
-
-@test gradtest(x -> x', rand(5))
-
-@test gradtest(vcat, rand(5), rand(3))
-@test gradtest(vcat, rand(2,3), rand(3,3))
-
-@test gradtest(rand(5)) do x
-  y = x.^2
-  2y + x
-end
-
-end

test/utils.jl

@@ -1,9 +1,13 @@
-using Flux: throttle
+using Flux
+using Flux: throttle, nfan, glorot_uniform, glorot_normal, stack, unstack
+using StatsBase: var
+using Random
+using Test
 
 @testset "Throttle" begin
   @testset "default behaviour" begin
     a = []
-    f = throttle(()->push!(a, now()), 1, leading=true, trailing=false)
+    f = throttle(()->push!(a, time()), 1, leading=true, trailing=false)
     f()
     f()
     f()
@@ -13,7 +17,7 @@ using Flux: throttle
 
   @testset "leading behaviour" begin
     a = []
-    f = throttle(()->push!(a, now()), 1, leading=true, trailing=false)
+    f = throttle(()->push!(a, time()), 1, leading=true, trailing=false)
     f()
     @test length(a) == 1
     f()
@@ -25,7 +29,7 @@ using Flux: throttle
 
   @testset "trailing behaviour" begin
     a = []
-    f = throttle(()->push!(a, now()), 1, leading=false, trailing=true)
+    f = throttle(()->push!(a, time()), 1, leading=false, trailing=true)
     f()
     @test length(a) == 0
     f()
@@ -47,3 +51,70 @@ using Flux: throttle
     @test a == [1, 3]
   end
 end
+
+@testset "Initialization" begin
+  # Set random seed so that these tests don't fail randomly
+  Random.seed!(0)
+
+  @testset "Fan in/out" begin
+    @test nfan() == (1, 1) #For a constant
+    @test nfan(100) == (1, 100) #For vector
+    @test nfan(100, 200) == (200, 100) #For Dense layer
+    @test nfan(2, 30, 40) == (2 * 30, 2 * 40) #For 1D Conv layer
+    @test nfan(2, 3, 40, 50) == (2 * 3 * 40, 2 * 3 * 50) #For 2D Conv layer
+    @test nfan(2, 3, 4, 50, 60) == (2 * 3 * 4 * 50, 2 * 3 * 4 * 60) #For 3D Conv layer
+  end
+
+  @testset "glorot" begin
+    # glorot_uniform and glorot_normal should both yield a kernel with
+    # variance ≈ 2/(fan_in + fan_out)
+    for dims ∈ [(1000,), (100, 100), (100, 400), (2, 3, 32, 64), (2, 3, 4, 32, 64)]
+      for init ∈ [glorot_uniform, glorot_normal]
+        v = init(dims...)
+        fan_in, fan_out = nfan(dims...)
+        σ2 = 2 / (fan_in + fan_out)
+        @test 0.9σ2 < var(v) < 1.1σ2
+      end
+    end
+  end
+end
+
+@testset "Params" begin
+  m = Dense(10, 5)
+  @test size.(params(m)) == [(5, 10), (5,)]
+  m = RNN(10, 5)
+  @test size.(params(m)) == [(5, 10), (5, 5), (5,), (5,)]
+
+  # Layer duplicated in same chain, params just once pls.
+  c = Chain(m, m)
+  @test size.(params(c)) == [(5, 10), (5, 5), (5,), (5,)]
+
+  # Self-referential array. Just want params, no stack overflow pls.
+  r = Any[nothing,m]
+  r[1] = r
+  @test size.(params(r)) == [(5, 10), (5, 5), (5,), (5,)]
+end
+
+@testset "Basic Stacking" begin
+  x = randn(3,3)
+  stacked = stack([x, x], 2)
+  @test size(stacked) == (3,2,3)
+end
+
+@testset "Precision" begin
+  m = Chain(Dense(10, 5, relu), Dense(5, 2))
+  x = rand(10)
+  @test eltype(m[1].W) == Float32
+  @test eltype(m(x)) == Float32
+  @test eltype(f64(m)(x)) == Float64
+  @test eltype(f64(m)[1].W) == Float64
+  @test eltype(f32(f64(m))[1].W) == Float32
+end
+
+@testset "Stacking" begin
+  stacked_array=[ 8 9 3 5; 9 6 6 9; 9 1 7 2; 7 4 10 6 ]
+  unstacked_array=[[8, 9, 9, 7], [9, 6, 1, 4], [3, 6, 7, 10], [5, 9, 2, 6]]
+  @test unstack(stacked_array, 2) == unstacked_array
+  @test stack(unstacked_array, 2) == stacked_array
+  @test stack(unstack(stacked_array, 1), 1) == stacked_array
+end