From 1e27430eacbb9bbb8d22861a758799db82435552 Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Sun, 9 Feb 2020 19:42:19 +0000
Subject: [PATCH 01/26] added LeakyReLU as a Bijector

---
 src/bijectors/leaky_relu.jl  | 102 ++++++++++++++++++++++++++++
 src/interface.jl             |   1 +
 test/bijectors/leaky_relu.jl | 127 +++++++++++++++++++++++++++++++++++
 test/runtests.jl             |   3 +
 4 files changed, 233 insertions(+)
 create mode 100644 src/bijectors/leaky_relu.jl
 create mode 100644 test/bijectors/leaky_relu.jl

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
new file mode 100644
index 00000000..0a15c681
--- /dev/null
+++ b/src/bijectors/leaky_relu.jl
@@ -0,0 +1,102 @@
+"""
+    LeakyReLU{T, N}(α::T) <: Bijector{N}
+
+Defines the invertible mapping
+
+    x ↦ x if x ≥ 0 else αx
+
+where α > 0.
+"""
+struct LeakyReLU{T, N} <: Bijector{N}
+    α::T
+end
+
+
+LeakyReLU(α::T; dim::Val{N} = Val(0)) where {T<:Real, N} = LeakyReLU{T, N}(α)
+LeakyReLU(α::T; dim::Val{N} = Val(D)) where {D, T<:AbstractArray{<:Real, D}, N} = LeakyReLU{T, N}(α)
+
+# (N=0) Univariate case
+function (b::LeakyReLU{T1, 0})(x::T2) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+    return T(ifelse(x < zero(T2), b.α * x, x))
+end
+(b::LeakyReLU{<:Real, 0})(x::AbstractVector{<:Real}) = b.(x)
+
+function (ib::Inversed{<:LeakyReLU{T1}, 0})(y::T2) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+    return T(ifelse(y < zero(T2), inv(ib.orig.α) * y, y))
+end
+(ib::Inversed{<:LeakyReLU{<:Real}, 0})(y::AbstractVector{<:Real}) = ib.(y)
+
+function logabsdetjac(b::LeakyReLU{T1, 0}, x::T2) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+    J⁻¹ = T(ifelse(x < 0, b.α, one(T2)))
+
+    return log.(abs.(J⁻¹))
+end
+logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.(b, x)
+
+
+# We implement `forward` by hand since we can re-use the computation of
+# the Jacobian of the transformation. This will lead to faster sampling
+# when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
+function forward(b::LeakyReLU{T1, 0}, x::T2) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+    J = T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+    return (rv=J * x, logabsdetjac=log(abs(J)))
+end
+
+# Batched version
+function forward(b::LeakyReLU{T1, 0}, x::AbstractVector{T2}) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+    J = @. T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+    return (rv=J .* x, logabsdetjac=log.(abs.(J)))
+end
+
+# (N=1) Multivariate case, with univariate parameter `α`
+function (b::LeakyReLU{T1, 1})(x::AbstractVecOrMat{T2}) where {T1<:Real, T2}
+    # Note that this will do the correct thing even for `Tracker.jl` in the sense
+    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
+    T = promote_type(T1, T2)
+    return @. T(ifelse(x < zero(T2), b.α * x, x))
+end
+
+function (ib::Inversed{<:LeakyReLU{T1}, 1})(y::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
+    # Note that this will do the correct thing even for `Tracker.jl` in the sense
+    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
+    T = promote_type(T1, T2)
+    return @. T(ifelse(y < zero(T2), inv(ib.orig.α) * y, y))
+end
+
+function logabsdetjac(b::LeakyReLU{T1, 1}, x::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
+    T = promote_type(T1, T2)
+
+    # Is really diagonal of jacobian
+    J⁻¹ = @. T(ifelse(x < 0, b.α, one(T2)))
+
+    if x isa AbstractVector
+        return sum(log.(abs.(J⁻¹)))
+    elseif x isa AbstractMatrix
+        return vec(sum(log.(abs.(J⁻¹)); dims = 1))  # sum along column
+    end
+end
+
+# We implement `forward` by hand since we can re-use the computation of
+# the Jacobian of the transformation. This will lead to faster sampling
+# when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
+function forward(b::LeakyReLU{T1, 1}, x::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
+    # Note that this will do the correct thing even for `Tracker.jl` in the sense
+    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
+    T = promote_type(T1, T2)
+
+    J = @. T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+
+    if x isa AbstractVector
+        logjac = sum(log.(abs.(J)))
+    elseif x isa AbstractMatrix
+        logjac = vec(sum(log.(abs.(J)); dims = 1))  # sum along column
+    end
+
+    y = J .* x
+    return (rv=y, logabsdetjac=logjac)
+end
diff --git a/src/interface.jl b/src/interface.jl
index 97998bfb..34117692 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -156,6 +156,7 @@ include("bijectors/distribution_bijector.jl")
 # Normalizing flow related
 include("bijectors/planar_layer.jl")
 include("bijectors/radial_layer.jl")
+include("bijectors/leaky_relu.jl")
 
 ##################
 # Other includes #
diff --git a/test/bijectors/leaky_relu.jl b/test/bijectors/leaky_relu.jl
new file mode 100644
index 00000000..cdf224b8
--- /dev/null
+++ b/test/bijectors/leaky_relu.jl
@@ -0,0 +1,127 @@
+using Test
+
+using Bijectors
+using Bijectors: LeakyReLU
+
+using LinearAlgebra
+using ForwardDiff
+
+true_logabsdetjac(b::Bijector{0}, x::Real) = (log ∘ abs)(ForwardDiff.derivative(b, x))
+true_logabsdetjac(b::Bijector{0}, x::AbstractVector) = (log ∘ abs).(ForwardDiff.derivative.(b, x))
+true_logabsdetjac(b::Bijector{1}, x::AbstractVector) = logabsdet(ForwardDiff.jacobian(b, x))[1]
+true_logabsdetjac(b::Bijector{1}, xs::AbstractMatrix) = mapreduce(z -> true_logabsdetjac(b, z), vcat, eachcol(xs))
+
+@testset "0-dim parameter, 0-dim input" begin
+    b = LeakyReLU(0.1; dim=Val(0))
+    x = 1.
+    @test inv(b)(b(x)) == x
+    @test inv(b)(b(-x)) == -x
+
+    # Mixing of types
+    # 1. Changes in input-type
+    @assert eltype(b(Float32(1.))) == Float64
+    @assert eltype(b(Float64(1.))) == Float64
+
+    # 2. Changes in parameter-type
+    b = LeakyReLU(Float32(0.1); dim=Val(0))
+    @assert eltype(b(Float32(1.))) == Float32
+    @assert eltype(b(Float64(1.))) == Float64
+
+    # logabsdetjac
+    @test logabsdetjac(b, x) == true_logabsdetjac(b, x)
+    @test logabsdetjac(b, Float32(x)) == true_logabsdetjac(b, x)
+
+    # Batch
+    xs = randn(10)
+    @test logabsdetjac(b, xs) == true_logabsdetjac(b, xs)
+    @test logabsdetjac(b, Float32.(x)) == true_logabsdetjac(b, Float32.(x))
+
+    @test logabsdetjac(b, -xs) == true_logabsdetjac(b, -xs)
+    @test logabsdetjac(b, -Float32.(xs)) == true_logabsdetjac(b, -Float32.(xs))
+
+    # Forward
+    f = forward(b, xs)
+    @test f.logabsdetjac ≈ logabsdetjac(b, xs)
+    @test f.rv ≈ b(xs)
+
+    f = forward(b, Float32.(xs))
+    @test f.logabsdetjac == logabsdetjac(b, Float32.(xs))
+    @test f.rv ≈ b(Float32.(xs))
+end
+
+@testset "0-dim parameter, 1-dim input" begin
+    d = 2
+
+    b = LeakyReLU(0.1; dim=Val(1))
+    x = ones(d)
+    @test inv(b)(b(x)) == x
+    @test inv(b)(b(-x)) == -x
+
+    # Batch
+    xs = randn(d, 10)
+    @test logabsdetjac(b, xs) == true_logabsdetjac(b, xs)
+    @test logabsdetjac(b, Float32.(x)) == true_logabsdetjac(b, Float32.(x))
+
+    @test logabsdetjac(b, -xs) == true_logabsdetjac(b, -xs)
+    @test logabsdetjac(b, -Float32.(xs)) == true_logabsdetjac(b, -Float32.(xs))
+
+    # Forward
+    f = forward(b, xs)
+    @test f.logabsdetjac ≈ logabsdetjac(b, xs)
+    @test f.rv ≈ b(xs)
+
+    f = forward(b, Float32.(xs))
+    @test f.logabsdetjac == logabsdetjac(b, Float32.(xs))
+    @test f.rv ≈ b(Float32.(xs))
+
+    # Mixing of types
+    # 1. Changes in input-type
+    @assert eltype(b(ones(Float32, 2))) == Float64
+    @assert eltype(b(ones(Float64, 2))) == Float64
+
+    # 2. Changes in parameter-type
+    b = LeakyReLU(Float32(0.1); dim=Val(1))
+    @assert eltype(b(ones(Float32, 2))) == Float32
+    @assert eltype(b(ones(Float64, 2))) == Float64
+
+    # grad_test_f(α, x) = (sum ∘ LeakyReLU(α; dim=Val(1)))(x)
+    # x = ones(2)
+    # α = [0.1]
+
+    # # Tracker.jl should have a gradient-type that is the same as the INPUT
+    # using Tracker
+    # res = Tracker.data(Tracker.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))[1])
+    # @test eltype(res) == Float32
+    # res = Tracker.data(Tracker.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))[1])
+    # @test eltype(res) == Float64
+
+    # res = Tracker.data(Tracker.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))[1])
+    # @test eltype(res) == Float64
+    # res = Tracker.data(Tracker.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))[1])
+    # @test eltype(res) == Float64
+
+    # x = Tracker.param(ones(Float32, 2))
+    # @test b(x) isa TrackedArray
+
+    # using Zygote
+    # res = Zygote.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))[1]
+    # @test eltype(res) == Float32
+    # res = Zygote.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))[1]
+    # @test eltype(res) == Float64
+
+    # res = Zygote.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))[1]
+    # @test eltype(res) == Float64
+    # res = Zygote.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))[1]
+    # @test eltype(res) == Float64
+
+    # using ForwardDiff
+    # res = ForwardDiff.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))
+    # @test eltype(res) == Float64
+    # res = ForwardDiff.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))
+    # @test eltype(res) == Float64
+
+    # res = ForwardDiff.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))
+    # @test eltype(res) == Float64
+    # res = ForwardDiff.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))
+    # @test eltype(res) == Float64
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 27df423d..7d612e1f 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -7,3 +7,6 @@ Random.seed!(123456)
     include("interface.jl")
 end
 include("transform.jl")
+@testset "Leaky ReLU" begin
+    include("bijectors/leaky_relu.jl")
+end

From 180a527502fff76c17c0138e38991cffddcaef9f Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Tue, 11 Feb 2020 13:44:39 +0000
Subject: [PATCH 02/26] added Compat.jl

---
 Project.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Project.toml b/Project.toml
index a617ffc7..1231be81 100644
--- a/Project.toml
+++ b/Project.toml
@@ -3,6 +3,7 @@ uuid = "76274a88-744f-5084-9051-94815aaf08c4"
 version = "0.5.3"
 
 [deps]
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
@@ -14,6 +15,7 @@ Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [compat]
+Compat = "3"
 Distributions = "0.21.11, 0.22"
 ForwardDiff = "0.10.3"
 MappedArrays = "0.2.2"

From 0f5200acf3c602108fac08fba14571407421b308 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Wed, 12 Feb 2020 09:38:03 +0000
Subject: [PATCH 03/26] forgot to use Compat.jl

---
 test/bijectors/leaky_relu.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/bijectors/leaky_relu.jl b/test/bijectors/leaky_relu.jl
index cdf224b8..f2ca4bc8 100644
--- a/test/bijectors/leaky_relu.jl
+++ b/test/bijectors/leaky_relu.jl
@@ -1,5 +1,7 @@
 using Test
 
+using Compat
+
 using Bijectors
 using Bijectors: LeakyReLU
 

From 533b86ea018c3c883622f19f32d2073d4fc6d433 Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Thu, 10 Sep 2020 07:48:35 +0200
Subject: [PATCH 04/26] added some tests for LeakyReLU

---
 test/interface.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/interface.jl b/test/interface.jl
index 19dd8c08..0bcb72b2 100644
--- a/test/interface.jl
+++ b/test/interface.jl
@@ -7,7 +7,7 @@ using Tracker
 using DistributionsAD
 
 using Bijectors
-using Bijectors: Log, Exp, Shift, Scale, Logit, SimplexBijector, PDBijector, Permute, PlanarLayer, RadialLayer, Stacked, TruncatedBijector, ADBijector
+using Bijectors: Log, Exp, Shift, Scale, Logit, SimplexBijector, PDBijector, Permute, PlanarLayer, RadialLayer, Stacked, TruncatedBijector, ADBijector, LeakyReLU
 
 Random.seed!(123)
 
@@ -161,6 +161,7 @@ end
         (Stacked((Exp{1}(), SimplexBijector()), [1:1, 2:3]),
          mapslices(z -> normalize(z, 1), rand(3, 2); dims = 1)),
         (LeakyReLU(0.1), randn(3)),
+        (LeakyReLU(Float32(0.1)), randn(3)),
         (LeakyReLU(0.1; dim = Val(1)), randn(2, 3))
     ]
 
@@ -174,7 +175,6 @@ end
             x = D == 0 ? xs[1] : xs[:, 1]
 
             y = @inferred b(x)
-
             ys = @inferred b(xs)
 
             # Computations which do not have closed-form implementations are not necessarily

From fd2e33be2a931077ea5988ed9c2935ead893d3d7 Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Thu, 10 Sep 2020 07:48:51 +0200
Subject: [PATCH 05/26] using masks rather than ifelse for type-stability and
 simplicity

---
 src/bijectors/leaky_relu.jl | 77 +++++++++++++++++--------------------
 1 file changed, 35 insertions(+), 42 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 0a15c681..8c612bc3 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -11,28 +11,28 @@ struct LeakyReLU{T, N} <: Bijector{N}
     α::T
 end
 
-
 LeakyReLU(α::T; dim::Val{N} = Val(0)) where {T<:Real, N} = LeakyReLU{T, N}(α)
 LeakyReLU(α::T; dim::Val{N} = Val(D)) where {D, T<:AbstractArray{<:Real, D}, N} = LeakyReLU{T, N}(α)
 
+up1(b::LeakyReLU{T, N}) where {T, N} = LeakyReLU{T, N + 1}(b.α)
+
 # (N=0) Univariate case
-function (b::LeakyReLU{T1, 0})(x::T2) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-    return T(ifelse(x < zero(T2), b.α * x, x))
+function (b::LeakyReLU{<:Any, 0})(x::Real)
+    mask = x < zero(x)
+    return mask * b.α * x + (1 - mask) * x
 end
-(b::LeakyReLU{<:Real, 0})(x::AbstractVector{<:Real}) = b.(x)
+(b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = b.(x)
 
-function (ib::Inversed{<:LeakyReLU{T1}, 0})(y::T2) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-    return T(ifelse(y < zero(T2), inv(ib.orig.α) * y, y))
+function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
+    mask = y < zero(y)
+    return mask * (x / b.α) + (1 - mask) * x
 end
-(ib::Inversed{<:LeakyReLU{<:Real}, 0})(y::AbstractVector{<:Real}) = ib.(y)
-
-function logabsdetjac(b::LeakyReLU{T1, 0}, x::T2) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-    J⁻¹ = T(ifelse(x < 0, b.α, one(T2)))
+(ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = ib.(y)
 
-    return log.(abs.(J⁻¹))
+function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
+    mask = x < zero(x)
+    J = mask * b.α + (1 - mask) * one(x)
+    return log.(abs.(J))
 end
 logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.(b, x)
 
@@ -40,56 +40,49 @@ logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.
 # We implement `forward` by hand since we can re-use the computation of
 # the Jacobian of the transformation. This will lead to faster sampling
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
-function forward(b::LeakyReLU{T1, 0}, x::T2) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-    J = T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+function forward(b::LeakyReLU{<:Any, 0}, x::Real)
+    mask = x < zero(x)
+    J = mask * b.α + (1 - mask) * one(x)
     return (rv=J * x, logabsdetjac=log(abs(J)))
 end
 
 # Batched version
-function forward(b::LeakyReLU{T1, 0}, x::AbstractVector{T2}) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-    J = @. T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
+    mask = x .< zero(eltype(x))
+    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
 end
 
 # (N=1) Multivariate case, with univariate parameter `α`
-function (b::LeakyReLU{T1, 1})(x::AbstractVecOrMat{T2}) where {T1<:Real, T2}
-    # Note that this will do the correct thing even for `Tracker.jl` in the sense
-    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
-    T = promote_type(T1, T2)
-    return @. T(ifelse(x < zero(T2), b.α * x, x))
+function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
+    mask = x .< zero(eltype(x))
+    return mask .* b.α .* x .+ (1 .- mask) .* x
 end
 
-function (ib::Inversed{<:LeakyReLU{T1}, 1})(y::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
-    # Note that this will do the correct thing even for `Tracker.jl` in the sense
-    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
-    T = promote_type(T1, T2)
-    return @. T(ifelse(y < zero(T2), inv(ib.orig.α) * y, y))
+function (ib::Inverse{<:LeakyReLU, 1})(y::AbstractVecOrMat)
+    mask = x .< zero(eltype(y))
+    return mask .* (y ./ ib.orig.α) .+ (1 .- mask) .* y
 end
 
-function logabsdetjac(b::LeakyReLU{T1, 1}, x::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
-    T = promote_type(T1, T2)
-
+function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    J⁻¹ = @. T(ifelse(x < 0, b.α, one(T2)))
+    mask = x .< zero(eltype(x))
+    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
 
     if x isa AbstractVector
-        return sum(log.(abs.(J⁻¹)))
+        return sum(log.(abs.(J)))
     elseif x isa AbstractMatrix
-        return vec(sum(log.(abs.(J⁻¹)); dims = 1))  # sum along column
+        return vec(sum(log.(abs.(J)); dims = 1))  # sum along column
     end
 end
 
 # We implement `forward` by hand since we can re-use the computation of
 # the Jacobian of the transformation. This will lead to faster sampling
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
-function forward(b::LeakyReLU{T1, 1}, x::AbstractVecOrMat{T2}) where {T1<:Real, T2<:Real}
-    # Note that this will do the correct thing even for `Tracker.jl` in the sense
-    # that the resulting array will be `TrackedArray` rather than `Array{<:TrackedReal}`.
-    T = promote_type(T1, T2)
-
-    J = @. T(ifelse(x < 0, b.α, one(T2))) # <= is really diagonal of jacobian
+function forward(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
+    # Is really diagonal of jacobian
+    mask = x .< zero(eltype(x))
+    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
 
     if x isa AbstractVector
         logjac = sum(log.(abs.(J)))

From 8b69cabd244c9589063597d81713a75687fc868d Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Thu, 10 Sep 2020 07:57:20 +0200
Subject: [PATCH 06/26] removed some redundant comments

---
 test/bijectors/leaky_relu.jl | 41 ------------------------------------
 1 file changed, 41 deletions(-)

diff --git a/test/bijectors/leaky_relu.jl b/test/bijectors/leaky_relu.jl
index f2ca4bc8..20aa8ffc 100644
--- a/test/bijectors/leaky_relu.jl
+++ b/test/bijectors/leaky_relu.jl
@@ -85,45 +85,4 @@ end
     b = LeakyReLU(Float32(0.1); dim=Val(1))
     @assert eltype(b(ones(Float32, 2))) == Float32
     @assert eltype(b(ones(Float64, 2))) == Float64
-
-    # grad_test_f(α, x) = (sum ∘ LeakyReLU(α; dim=Val(1)))(x)
-    # x = ones(2)
-    # α = [0.1]
-
-    # # Tracker.jl should have a gradient-type that is the same as the INPUT
-    # using Tracker
-    # res = Tracker.data(Tracker.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))[1])
-    # @test eltype(res) == Float32
-    # res = Tracker.data(Tracker.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))[1])
-    # @test eltype(res) == Float64
-
-    # res = Tracker.data(Tracker.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))[1])
-    # @test eltype(res) == Float64
-    # res = Tracker.data(Tracker.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))[1])
-    # @test eltype(res) == Float64
-
-    # x = Tracker.param(ones(Float32, 2))
-    # @test b(x) isa TrackedArray
-
-    # using Zygote
-    # res = Zygote.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))[1]
-    # @test eltype(res) == Float32
-    # res = Zygote.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))[1]
-    # @test eltype(res) == Float64
-
-    # res = Zygote.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))[1]
-    # @test eltype(res) == Float64
-    # res = Zygote.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))[1]
-    # @test eltype(res) == Float64
-
-    # using ForwardDiff
-    # res = ForwardDiff.gradient(x -> grad_test_f(Float64(α[1]), x), Float32.(x))
-    # @test eltype(res) == Float64
-    # res = ForwardDiff.gradient(x -> grad_test_f(Float32(α[1]), x), Float64.(x))
-    # @test eltype(res) == Float64
-
-    # res = ForwardDiff.gradient(α -> grad_test_f(α[1], Float64.(x)), Float32.(α))
-    # @test eltype(res) == Float64
-    # res = ForwardDiff.gradient(α -> grad_test_f(α[1], Float32.(x)), Float64.(α))
-    # @test eltype(res) == Float64
 end

From ba299ff01ecc88b0404a070aeb078a98888df89c Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 07:04:54 +0200
Subject: [PATCH 07/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 8c612bc3..c0782a1a 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -25,7 +25,7 @@ end
 
 function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
     mask = y < zero(y)
-    return mask * (x / b.α) + (1 - mask) * x
+    return mask * (x / ib.orig.α) + (1 - mask) * x
 end
 (ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = ib.(y)
 

From 9b4d0d6bee762828d4e51d6ea48d791d5adee0ab Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 09:43:04 +0200
Subject: [PATCH 08/26] removed unnecessary broadcasting and useless import

---
 src/bijectors/leaky_relu.jl  | 2 +-
 test/bijectors/leaky_relu.jl | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 8c612bc3..eff20614 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -32,7 +32,7 @@ end
 function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)
     J = mask * b.α + (1 - mask) * one(x)
-    return log.(abs.(J))
+    return log(abs(J))
 end
 logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.(b, x)
 
diff --git a/test/bijectors/leaky_relu.jl b/test/bijectors/leaky_relu.jl
index 20aa8ffc..63ba8c18 100644
--- a/test/bijectors/leaky_relu.jl
+++ b/test/bijectors/leaky_relu.jl
@@ -1,7 +1,5 @@
 using Test
 
-using Compat
-
 using Bijectors
 using Bijectors: LeakyReLU
 

From 7771155094b52c3a7e970e05caadcb685082557c Mon Sep 17 00:00:00 2001
From: "tor.erlend95@gmail.com" <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 15:47:47 +0200
Subject: [PATCH 09/26] fixed a typo

---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 713d003f..b4acfd63 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -25,7 +25,7 @@ end
 
 function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
     mask = y < zero(y)
-    return mask * (x / ib.orig.α) + (1 - mask) * x
+    return mask * (y / ib.orig.α) + (1 - mask) * y
 end
 (ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = ib.(y)
 

From 6e9d9f7fc794441a85de66ac2885c2bc2bb593ad Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:15:10 +0200
Subject: [PATCH 10/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index b4acfd63..b581156d 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -49,7 +49,7 @@ end
 # Batched version
 function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
     mask = x .< zero(eltype(x))
-    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
+    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
 end
 

From 18e30f9a786c496ab7f0031fdec618a88b341b8a Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:15:49 +0200
Subject: [PATCH 11/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index b581156d..45e8fe2c 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -19,7 +19,7 @@ up1(b::LeakyReLU{T, N}) where {T, N} = LeakyReLU{T, N + 1}(b.α)
 # (N=0) Univariate case
 function (b::LeakyReLU{<:Any, 0})(x::Real)
     mask = x < zero(x)
-    return mask * b.α * x + (1 - mask) * x
+    return mask * b.α * x + !mask * x
 end
 (b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = b.(x)
 

From c5b058bf76a9d42e9981e41301bf82a3493ac528 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:16:01 +0200
Subject: [PATCH 12/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 45e8fe2c..3571c686 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -21,7 +21,7 @@ function (b::LeakyReLU{<:Any, 0})(x::Real)
     mask = x < zero(x)
     return mask * b.α * x + !mask * x
 end
-(b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = b.(x)
+(b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = map(b, x)
 
 function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
     mask = y < zero(y)

From c72fd1152f12c2bf6072f0a0e28a619bb74572fa Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:16:09 +0200
Subject: [PATCH 13/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 3571c686..70b52abd 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -25,7 +25,7 @@ end
 
 function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
     mask = y < zero(y)
-    return mask * (y / ib.orig.α) + (1 - mask) * y
+    return mask * (y / ib.orig.α) + !mask * y
 end
 (ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = ib.(y)
 

From 68d35188335c3f85ded30752f87cb6c23126f555 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:16:17 +0200
Subject: [PATCH 14/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 70b52abd..0d26f877 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -27,7 +27,7 @@ function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
     mask = y < zero(y)
     return mask * (y / ib.orig.α) + !mask * y
 end
-(ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = ib.(y)
+(ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = map(ib, y)
 
 function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)

From c14a5db88402d85ad9d71551a17341ee3ed37f45 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:16:29 +0200
Subject: [PATCH 15/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 0d26f877..73816066 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -42,7 +42,7 @@ logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
 function forward(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)
-    J = mask * b.α + (1 - mask) * one(x)
+    J = mask * b.α + !mask * one(x)
     return (rv=J * x, logabsdetjac=log(abs(J)))
 end
 

From 4a21b73ab2a2b132972f44b2eba8c2c53180ea3b Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:16:56 +0200
Subject: [PATCH 16/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 73816066..680e0a85 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -55,8 +55,7 @@ end
 
 # (N=1) Multivariate case, with univariate parameter `α`
 function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
-    mask = x .< zero(eltype(x))
-    return mask .* b.α .* x .+ (1 .- mask) .* x
+    return @. (x < zero(x)) * b.α * x + (x > zero(x)) * x
 end
 
 function (ib::Inverse{<:LeakyReLU, 1})(y::AbstractVecOrMat)

From b4119fda4a3952313f26fc48793f602e5015b593 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:17:03 +0200
Subject: [PATCH 17/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 680e0a85..5d138228 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -65,8 +65,7 @@ end
 
 function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    mask = x .< zero(eltype(x))
-    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
+    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
 
     if x isa AbstractVector
         return sum(log.(abs.(J)))

From 1b97a0f98c2a1d0610acc6ccbafb284bc5547314 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:17:11 +0200
Subject: [PATCH 18/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 5d138228..5dd91545 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -79,8 +79,7 @@ end
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
 function forward(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    mask = x .< zero(eltype(x))
-    J = mask .* b.α .+ (1 .- mask) .* one(eltype(x))
+    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
 
     if x isa AbstractVector
         logjac = sum(log.(abs.(J)))

From 1ebd9ba01084fa2c9898d32ce5e53e7fdacfbd59 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:19:03 +0200
Subject: [PATCH 19/26] Apply suggestions from code review

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 5dd91545..088c6894 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -59,8 +59,7 @@ function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
 end
 
 function (ib::Inverse{<:LeakyReLU, 1})(y::AbstractVecOrMat)
-    mask = x .< zero(eltype(y))
-    return mask .* (y ./ ib.orig.α) .+ (1 .- mask) .* y
+    return @. (y < zero(y)) * y / ib.orig.α + (y > zero(x)) * y
 end
 
 function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)

From 07cc6328cd9b6543b7cca00d90c638b035ef951c Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 11 Sep 2020 16:19:29 +0200
Subject: [PATCH 20/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 088c6894..3db1f82d 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -34,7 +34,7 @@ function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     J = mask * b.α + (1 - mask) * one(x)
     return log(abs(J))
 end
-logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = logabsdetjac.(b, x)
+logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = map(x -> logabsdetjac(b, x), x)
 
 
 # We implement `forward` by hand since we can re-use the computation of

From f2f167e80c6845b716d6a68b6d85dccf1b656d21 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 13:07:21 +0200
Subject: [PATCH 21/26] Apply suggestions from code review

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 3db1f82d..49fc703b 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -23,12 +23,7 @@ function (b::LeakyReLU{<:Any, 0})(x::Real)
 end
 (b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = map(b, x)
 
-function (ib::Inverse{<:LeakyReLU, 0})(y::Real)
-    mask = y < zero(y)
-    return mask * (y / ib.orig.α) + !mask * y
-end
-(ib::Inverse{<:LeakyReLU{<:Any}, 0})(y::AbstractVector{<:Real}) = map(ib, y)
-
+Base.inv(b::LeakyReLU) = LeakyReLU(inv.(b.a))
 function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)
     J = mask * b.α + (1 - mask) * one(x)
@@ -48,7 +43,6 @@ end
 
 # Batched version
 function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
-    mask = x .< zero(eltype(x))
     J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
 end
@@ -58,10 +52,6 @@ function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
     return @. (x < zero(x)) * b.α * x + (x > zero(x)) * x
 end
 
-function (ib::Inverse{<:LeakyReLU, 1})(y::AbstractVecOrMat)
-    return @. (y < zero(y)) * y / ib.orig.α + (y > zero(x)) * y
-end
-
 function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
     J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)

From 003dfb6f5888a753102069c9063cae1d8d10be52 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 16:49:59 +0200
Subject: [PATCH 22/26] Apply suggestions from code review

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 49fc703b..9fd06eff 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -47,7 +47,7 @@ function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
 end
 
-# (N=1) Multivariate case, with univariate parameter `α`
+# (N=1) Multivariate case
 function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
     return @. (x < zero(x)) * b.α * x + (x > zero(x)) * x
 end

From e17d5d74651520887cabbe5ca1801ddf5f6a7831 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 16:51:38 +0200
Subject: [PATCH 23/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 9fd06eff..63b7845b 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -23,7 +23,7 @@ function (b::LeakyReLU{<:Any, 0})(x::Real)
 end
 (b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = map(b, x)
 
-Base.inv(b::LeakyReLU) = LeakyReLU(inv.(b.a))
+Base.inv(b::LeakyReLU) = LeakyReLU(inv.(b.α))
 function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)
     J = mask * b.α + (1 - mask) * one(x)

From a83cb7b33eac6b28dbccb7d86e480e973b5010e4 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 19:22:17 +0200
Subject: [PATCH 24/26] Update src/bijectors/leaky_relu.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/bijectors/leaky_relu.jl | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 63b7845b..be94f247 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -23,7 +23,11 @@ function (b::LeakyReLU{<:Any, 0})(x::Real)
 end
 (b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = map(b, x)
 
-Base.inv(b::LeakyReLU) = LeakyReLU(inv.(b.α))
+function Base.inv(b::LeakyReLU{<:Any,N}) where N
+    invα = inv.(b.α)
+    return LeakyReLU{typeof(invα),N}(invα)
+end
+
 function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
     mask = x < zero(x)
     J = mask * b.α + (1 - mask) * one(x)

From 28158ae939145b28270d7443fe327344f4284b9f Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 21:44:48 +0200
Subject: [PATCH 25/26] Apply suggestions from code review

---
 src/bijectors/leaky_relu.jl | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index be94f247..56fcfd44 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -47,18 +47,24 @@ end
 
 # Batched version
 function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
-    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
+    J = let z = zero(x), o = one(x)
+        @. (x < z) * b.α + (x > z) * o
+    end
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
 end
 
 # (N=1) Multivariate case
 function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
-    return @. (x < zero(x)) * b.α * x + (x > zero(x)) * x
+    return let z = zero(x)
+        @. (x < z) * b.α * x + (x > z) * x
+    end
 end
 
 function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
+    J = let z = zero(x), o = one(x)
+        @. (x < z) * b.α + (x > z) * o
+    end
 
     if x isa AbstractVector
         return sum(log.(abs.(J)))
@@ -72,7 +78,9 @@ end
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
 function forward(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    J = @. (x < zero(x)) * b.α + (x > zero(x)) * one(x)
+    J = let z = zero(x), o = one(x)
+        @. (x < z) * b.α + (x > z) * o
+    end
 
     if x isa AbstractVector
         logjac = sum(log.(abs.(J)))

From 60775786ce4fe8a901651a5660a8ef3efe257259 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Sat, 12 Sep 2020 21:53:43 +0200
Subject: [PATCH 26/26] Apply suggestions from code review

---
 src/bijectors/leaky_relu.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/bijectors/leaky_relu.jl b/src/bijectors/leaky_relu.jl
index 56fcfd44..ffca00c7 100644
--- a/src/bijectors/leaky_relu.jl
+++ b/src/bijectors/leaky_relu.jl
@@ -47,7 +47,7 @@ end
 
 # Batched version
 function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
-    J = let z = zero(x), o = one(x)
+    J = let T = eltype(x), z = zero(T), o = one(T)
         @. (x < z) * b.α + (x > z) * o
     end
     return (rv=J .* x, logabsdetjac=log.(abs.(J)))
@@ -55,14 +55,14 @@ end
 
 # (N=1) Multivariate case
 function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
-    return let z = zero(x)
+    return let z = zero(eltype(x))
         @. (x < z) * b.α * x + (x > z) * x
     end
 end
 
 function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    J = let z = zero(x), o = one(x)
+    J = let T = eltype(x), z = zero(T), o = one(T)
         @. (x < z) * b.α + (x > z) * o
     end
 
@@ -78,7 +78,7 @@ end
 # when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
 function forward(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
     # Is really diagonal of jacobian
-    J = let z = zero(x), o = one(x)
+    J = let T = eltype(x), z = zero(T), o = one(T)
         @. (x < z) * b.α + (x > z) * o
     end