From 2c8f70c6fb8a942beb5f04f6d753ee2725d3e00f Mon Sep 17 00:00:00 2001
From: Piotr Tarasiewicz <taras.piotr@gmail.com>
Date: Sun, 25 Oct 2020 11:11:26 +0000
Subject: [PATCH] tensor

---
 src/DeepJulia.jl           |  9 +++++++--
 src/device.jl              |  4 ++++
 src/modules.jl             |  4 ++--
 src/optim.jl               |  2 +-
 src/{cuda.jl => tensor.jl} | 11 +++--------
 src/variable.jl            |  4 ++--
 6 files changed, 19 insertions(+), 15 deletions(-)
 create mode 100644 src/device.jl
 rename src/{cuda.jl => tensor.jl} (51%)

diff --git a/src/DeepJulia.jl b/src/DeepJulia.jl
index d553163..c8cf4eb 100644
--- a/src/DeepJulia.jl
+++ b/src/DeepJulia.jl
@@ -3,8 +3,12 @@ module DeepJulia
 export Device,
 gpu,
 cpu,
-ArrayOrCuArray,
+
+# tensor
+
+Tensor,
 to,
+device,
 
 # variable
 
@@ -44,7 +48,8 @@ batchify,
 FashionMNIST,
 shuffle!
 
-include("cuda.jl")
+include("device.jl")
+include("tensor.jl")
 include("variable.jl")
 include("loss.jl")
 include("modules.jl")
diff --git a/src/device.jl b/src/device.jl
new file mode 100644
index 0000000..99c4e14
--- /dev/null
+++ b/src/device.jl
@@ -0,0 +1,4 @@
+@enum Device begin
+    cpu
+    gpu
+end
diff --git a/src/modules.jl b/src/modules.jl
index a61cca4..845fb49 100644
--- a/src/modules.jl
+++ b/src/modules.jl
@@ -12,9 +12,9 @@ mutable struct LinearLayer <: NNModule
     W::Variable
     b::Variable
     
-    input::ArrayOrCuArray
+    input::Tensor
     
-    LinearLayer(W::ArrayOrCuArray, b::ArrayOrCuArray) = new(
+    LinearLayer(W::Tensor, b::Tensor) = new(
         Variable(W),
         Variable(b),
         Matrix{Real}(undef, 0, 0),
diff --git a/src/optim.jl b/src/optim.jl
index 4e5f851..e3bfcba 100644
--- a/src/optim.jl
+++ b/src/optim.jl
@@ -7,7 +7,7 @@ struct SGD <: Optimizer
     params::Vector{Variable}
     lr::Real
     momentum::Real
-    velocities::Vector{ArrayOrCuArray}
+    velocities::Vector{Tensor}
     
     SGD(params, lr) = new(params, lr, 0.0, Vector{Matrix{Real}}(undef, size(params, 1)))
     SGD(params, lr, momentum) = new(params, lr, momentum, [to(zeros(size(p.values)), device(p.values)) for p ∈ params])
diff --git a/src/cuda.jl b/src/tensor.jl
similarity index 51%
rename from src/cuda.jl
rename to src/tensor.jl
index 171dd06..30f5c24 100644
--- a/src/cuda.jl
+++ b/src/tensor.jl
@@ -1,13 +1,8 @@
 using CUDA: CuArray
 
-@enum Device begin
-    cpu
-    gpu
-end
-
-ArrayOrCuArray = Union{Array,CuArray}
+Tensor = Union{Array,CuArray}
 
-function to(A::ArrayOrCuArray, device::Device)
+function to(A::Tensor, device::Device)
     if device == cpu
         return Array(A)
     elseif device == gpu
@@ -17,4 +12,4 @@ function to(A::ArrayOrCuArray, device::Device)
     end
 end
 
-device(A::ArrayOrCuArray) = isa(A, CuArray) ? gpu : cpu
+device(A::Tensor) = isa(A, CuArray) ? gpu : cpu
diff --git a/src/variable.jl b/src/variable.jl
index 03ba2c6..99b33d1 100644
--- a/src/variable.jl
+++ b/src/variable.jl
@@ -1,6 +1,6 @@
 mutable struct Variable
-    values::ArrayOrCuArray
-    grad::ArrayOrCuArray
+    values::Tensor
+    grad::Tensor
     
     Variable(values) = new(values, to(zeros(size(values)), device(values)))
     Variable(values, grad) = new(values, grad)