From 36c5256bcb80a92853cfe6231287e9d9c8dca88e Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Mon, 27 May 2019 15:14:02 -0700
Subject: [PATCH 1/4] Make jacobian work with CuArrays

---
 src/back.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/back.jl b/src/back.jl
index e638e1ad..329e5fcc 100644
--- a/src/back.jl
+++ b/src/back.jl
@@ -173,7 +173,11 @@ Calculate the output jacobian `J = d/dx m(x)` such that each row `i` of `J` corr
 """
 function jacobian(f, x::AbstractVector)
   y::AbstractVector, back = forward(f, x)
-  ȳ(i) = [i == j for j = 1:length(y)]
+  function ȳ(i)
+    δ = fill!(similar(y, Bool), false)
+    δ[i] = true
+    return δ
+  end
   vcat([transpose(back(ȳ(i))[1]) for i = 1:length(y)]...)
 end
 

From aa4e64a927104104a15ac506e6f626ef0da75d8d Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Tue, 28 May 2019 20:09:03 -0700
Subject: [PATCH 2/4] Use broadcasting to construct the Kronecker's delta

---
 src/back.jl | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/back.jl b/src/back.jl
index 329e5fcc..db9669f5 100644
--- a/src/back.jl
+++ b/src/back.jl
@@ -173,11 +173,8 @@ Calculate the output jacobian `J = d/dx m(x)` such that each row `i` of `J` corr
 """
 function jacobian(f, x::AbstractVector)
   y::AbstractVector, back = forward(f, x)
-  function ȳ(i)
-    δ = fill!(similar(y, Bool), false)
-    δ[i] = true
-    return δ
-  end
+  # Using broadcasting so that output of `ȳ` is a GPU array if `y` is so:
+  ȳ(i) = ((j, _) -> i == j).(1:length(y), y)
   vcat([transpose(back(ȳ(i))[1]) for i = 1:length(y)]...)
 end
 

From 9c2c05ec6d785996d83f7f59e3209461fcc36c01 Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Wed, 29 May 2019 22:10:58 -0700
Subject: [PATCH 3/4] =?UTF-8?q?Make=20sure=20that=20y=CC=84's=20output=20i?=
 =?UTF-8?q?s=20a=20compatible=20float=20type?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/back.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/back.jl b/src/back.jl
index db9669f5..c88a5100 100644
--- a/src/back.jl
+++ b/src/back.jl
@@ -173,8 +173,9 @@ Calculate the output jacobian `J = d/dx m(x)` such that each row `i` of `J` corr
 """
 function jacobian(f, x::AbstractVector)
   y::AbstractVector, back = forward(f, x)
+  z = float(zero(eltype(data(y))))
   # Using broadcasting so that output of `ȳ` is a GPU array if `y` is so:
-  ȳ(i) = ((j, _) -> i == j).(1:length(y), y)
+  ȳ(i) = ((j, _) -> i == j).(1:length(y), y) .+ z
   vcat([transpose(back(ȳ(i))[1]) for i = 1:length(y)]...)
 end
 

From f6b0ad40a196059b50fd95045b97607ae82b0232 Mon Sep 17 00:00:00 2001
From: Takafumi Arakaki <aka.tkf@gmail.com>
Date: Wed, 29 May 2019 23:50:07 -0700
Subject: [PATCH 4/4] Use `fill!`/`similar`-based approach to reduce number of
 indexing

---
 src/back.jl | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/back.jl b/src/back.jl
index c88a5100..c34e2a73 100644
--- a/src/back.jl
+++ b/src/back.jl
@@ -173,9 +173,11 @@ Calculate the output jacobian `J = d/dx m(x)` such that each row `i` of `J` corr
 """
 function jacobian(f, x::AbstractVector)
   y::AbstractVector, back = forward(f, x)
-  z = float(zero(eltype(data(y))))
-  # Using broadcasting so that output of `ȳ` is a GPU array if `y` is so:
-  ȳ(i) = ((j, _) -> i == j).(1:length(y), y) .+ z
+  function ȳ(i)
+    δ = fill!(float(similar(data(y))), false)
+    δ[i] = true
+    return δ
+  end
   vcat([transpose(back(ȳ(i))[1]) for i = 1:length(y)]...)
 end