From 466b5c501ae175a6ce69a2b74fae1a7c1c876da2 Mon Sep 17 00:00:00 2001
From: Mike J Innes <mike.j.innes@gmail.com>
Date: Mon, 26 Feb 2018 23:10:59 +0000
Subject: [PATCH] cpu/gpu conveniences

---
 src/Flux.jl     |  2 +-
 src/treelike.jl | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/Flux.jl b/src/Flux.jl
index 2835f281..110fd966 100644
--- a/src/Flux.jl
+++ b/src/Flux.jl
@@ -10,7 +10,7 @@ using MacroTools: @forward
 export Chain, Dense, RNN, LSTM, GRU, Conv, Conv2D,
   Dropout, LayerNorm, BatchNorm,
   SGD, ADAM, Momentum, Nesterov, AMSGrad,
-  param, params, mapleaves
+  param, params, mapleaves, cpu, gpu
 
 @reexport using NNlib
 
diff --git a/src/treelike.jl b/src/treelike.jl
index 235aaf02..cc7d827a 100644
--- a/src/treelike.jl
+++ b/src/treelike.jl
@@ -40,3 +40,15 @@ function params(m)
 end
 
 params(m...) = params(m)
+
+# CPU/GPU movement conveniences
+
+cpu(x) = adapt(Array, x)
+
+default_adaptor = Array
+
+@require CuArrays begin
+  global default_adaptor = CuArray
+end
+
+gpu(x) = adapt(default_adaptor, x)