From 730ee838d0e72ee8716889617eec2ffec12ce69b Mon Sep 17 00:00:00 2001
From: nic-barbara <nbar5346@uni.sydney.edu.au>
Date: Wed, 12 Jul 2023 18:52:55 +1000
Subject: [PATCH 1/8] Bug still there when nu, ny = 1

---
 test/Wrappers/diff_ren.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/test/Wrappers/diff_ren.jl b/test/Wrappers/diff_ren.jl
index c7a649f..837e6f9 100644
--- a/test/Wrappers/diff_ren.jl
+++ b/test/Wrappers/diff_ren.jl
@@ -11,10 +11,9 @@ using Test
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny = 4, 5, 10, 3
+nu, nx, nv, ny = 1, 10, 0, 1
 γ = 10
-ren_ps = LipschitzRENParams{Float32}(nu, nx, nv, ny, γ)
-model = DiffREN(ren_ps)
+model_ps = LipschitzRENParams{Float32}(nu, nx, nv, ny, γ)
 
 # Dummy data
 us = randn(nu, batches)
@@ -22,16 +21,17 @@ ys = randn(ny, batches)
 data = [(us[:,k], ys[:,k]) for k in 1:batches]
 
 # Dummy loss function just for testing
-function loss(m, u, y)
+function loss(model_ps, u, y)
+    m = DiffREN(model_ps)
     x0 = init_states(m, size(u,2))
     x1, y1 = m(x0, u)
     return Flux.mse(y1, y) + sum(x1.^2)
 end
 
 # Check if parameters change after a Flux update
-ps1 = deepcopy(Flux.params(model))
-opt_state = Flux.setup(Adam(0.01), model)
-Flux.train!(loss, model, data, opt_state)
-ps2 = Flux.params(model)
+ps1 = deepcopy(Flux.params(model_ps))
+opt_state = Flux.setup(Adam(0.01), model_ps)
+Flux.train!(loss, model_ps, data, opt_state)
+ps2 = Flux.params(model_ps)
 
 @test !any(ps1 .≈ ps2)

From 69d08dfbf19734c8b5ab1cdb0c22cb695e4eae06 Mon Sep 17 00:00:00 2001
From: nic-barbara <nbar5346@uni.sydney.edu.au>
Date: Wed, 12 Jul 2023 19:42:27 +1000
Subject: [PATCH 2/8] Minor edits

---
 src/ParameterTypes/utils.jl | 1 +
 test/Wrappers/diff_ren.jl   | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/ParameterTypes/utils.jl b/src/ParameterTypes/utils.jl
index bb5bc13..139c79b 100644
--- a/src/ParameterTypes/utils.jl
+++ b/src/ParameterTypes/utils.jl
@@ -78,6 +78,7 @@ function hmatrix_to_explicit(ps::AbstractRENParams, H::AbstractMatrix{T}, D22::A
     B1 = E \ B1_imp
     B2 = E \ ps.direct.B2
 
+    # Current versions of Julia behave poorly when broadcasting over 0-dim arrays
     C1  = (nv == 0) ? zeros(T,0,nx) : broadcast(*, Λ_inv, C1_imp)
     D11 = (nv == 0) ? zeros(T,0,0)  : broadcast(*, Λ_inv, D11_imp)
     D12 = (nv == 0) ? zeros(T,0,nu) : broadcast(*, Λ_inv, ps.direct.D12)
diff --git a/test/Wrappers/diff_ren.jl b/test/Wrappers/diff_ren.jl
index 837e6f9..7fe2281 100644
--- a/test/Wrappers/diff_ren.jl
+++ b/test/Wrappers/diff_ren.jl
@@ -11,7 +11,7 @@ using Test
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny = 1, 10, 0, 1
+nu, nx, nv, ny = 4, 5, 10, 2
 γ = 10
 model_ps = LipschitzRENParams{Float32}(nu, nx, nv, ny, γ)
 

From 92743b74060d752b1580f1f8256ea142e31310a9 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Wed, 12 Jul 2023 23:44:19 +1000
Subject: [PATCH 3/8] Isolated the bug in an example

---
 test/Wrappers/diff_ren.jl | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/test/Wrappers/diff_ren.jl b/test/Wrappers/diff_ren.jl
index 7fe2281..aa4ebbd 100644
--- a/test/Wrappers/diff_ren.jl
+++ b/test/Wrappers/diff_ren.jl
@@ -5,15 +5,16 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-# include("../test_utils.jl")
+include("../test_utils.jl")
 
 """
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny = 4, 5, 10, 2
+nu, nx, nv, ny = 4, 5, 0, 2
 γ = 10
-model_ps = LipschitzRENParams{Float32}(nu, nx, nv, ny, γ)
+ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
+model = DiffREN(ren_ps)
 
 # Dummy data
 us = randn(nu, batches)
@@ -21,17 +22,22 @@ ys = randn(ny, batches)
 data = [(us[:,k], ys[:,k]) for k in 1:batches]
 
 # Dummy loss function just for testing
-function loss(model_ps, u, y)
-    m = DiffREN(model_ps)
+function loss(m, u, y)
     x0 = init_states(m, size(u,2))
     x1, y1 = m(x0, u)
     return Flux.mse(y1, y) + sum(x1.^2)
 end
 
-# Check if parameters change after a Flux update
-ps1 = deepcopy(Flux.params(model_ps))
-opt_state = Flux.setup(Adam(0.01), model_ps)
-Flux.train!(loss, model_ps, data, opt_state)
-ps2 = Flux.params(model_ps)
+# Debug batch updates
+opt_state = Flux.setup(Adam(0.01), model)
+gs = Flux.gradient(loss, model, us, ys)
+Flux.update!(opt_state, model, gs[1])
+println()
 
-@test !any(ps1 .≈ ps2)
+# # Check if parameters change after a Flux update
+# ps1 = deepcopy(Flux.params(model))
+# opt_state = Flux.setup(Adam(0.01), model)
+# Flux.train!(loss, model, data, opt_state)
+# ps2 = Flux.params(model)
+
+# @test !any(ps1 .≈ ps2)

From 0a344a40d49b89ab7738d3322d150d98a59548a4 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Thu, 13 Jul 2023 07:21:46 +1000
Subject: [PATCH 4/8] Reduced minimal example of bug

---
 test/Wrappers/diff_ren.jl | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/test/Wrappers/diff_ren.jl b/test/Wrappers/diff_ren.jl
index aa4ebbd..c589baf 100644
--- a/test/Wrappers/diff_ren.jl
+++ b/test/Wrappers/diff_ren.jl
@@ -5,21 +5,20 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-include("../test_utils.jl")
+# include("../test_utils.jl")
 
 """
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny = 4, 5, 0, 2
-γ = 10
+nu, nx, nv, ny, γ = 4, 5, 0, 2, 10
 ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
 model = DiffREN(ren_ps)
 
 # Dummy data
 us = randn(nu, batches)
 ys = randn(ny, batches)
-data = [(us[:,k], ys[:,k]) for k in 1:batches]
+data = [(us, ys)]
 
 # Dummy loss function just for testing
 function loss(m, u, y)
@@ -31,7 +30,6 @@ end
 # Debug batch updates
 opt_state = Flux.setup(Adam(0.01), model)
 gs = Flux.gradient(loss, model, us, ys)
-Flux.update!(opt_state, model, gs[1])
 println()
 
 # # Check if parameters change after a Flux update

From 0c649c261a15a14455a0e75b8dc5984e3b0ee585 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Thu, 13 Jul 2023 08:10:54 +1000
Subject: [PATCH 5/8] Created new test for this bug

---
 test/Wrappers/diff_lbdn.jl |  2 --
 test/Wrappers/diff_ren.jl  | 22 ++++++---------
 test/Wrappers/zero_dim.jl  | 56 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 16 deletions(-)
 create mode 100644 test/Wrappers/zero_dim.jl

diff --git a/test/Wrappers/diff_lbdn.jl b/test/Wrappers/diff_lbdn.jl
index 25ef858..f629bce 100644
--- a/test/Wrappers/diff_lbdn.jl
+++ b/test/Wrappers/diff_lbdn.jl
@@ -5,8 +5,6 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-# include("../test_utils.jl")
-
 """
 Test that backpropagation runs and parameters change
 """
diff --git a/test/Wrappers/diff_ren.jl b/test/Wrappers/diff_ren.jl
index c589baf..8357243 100644
--- a/test/Wrappers/diff_ren.jl
+++ b/test/Wrappers/diff_ren.jl
@@ -5,20 +5,19 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-# include("../test_utils.jl")
-
 """
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny, γ = 4, 5, 0, 2, 10
+nu, nx, nv, ny = 4, 5, 10, 2
+γ = 10
 ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
 model = DiffREN(ren_ps)
 
 # Dummy data
 us = randn(nu, batches)
 ys = randn(ny, batches)
-data = [(us, ys)]
+data = [(us[:,k], ys[:,k]) for k in 1:batches]
 
 # Dummy loss function just for testing
 function loss(m, u, y)
@@ -27,15 +26,10 @@ function loss(m, u, y)
     return Flux.mse(y1, y) + sum(x1.^2)
 end
 
-# Debug batch updates
+# Check if parameters change after a Flux update
+ps1 = deepcopy(Flux.params(model))
 opt_state = Flux.setup(Adam(0.01), model)
-gs = Flux.gradient(loss, model, us, ys)
-println()
-
-# # Check if parameters change after a Flux update
-# ps1 = deepcopy(Flux.params(model))
-# opt_state = Flux.setup(Adam(0.01), model)
-# Flux.train!(loss, model, data, opt_state)
-# ps2 = Flux.params(model)
+Flux.train!(loss, model, data, opt_state)
+ps2 = Flux.params(model)
 
-# @test !any(ps1 .≈ ps2)
+@test !any(ps1 .≈ ps2)
diff --git a/test/Wrappers/zero_dim.jl b/test/Wrappers/zero_dim.jl
new file mode 100644
index 0000000..1a2d4d0
--- /dev/null
+++ b/test/Wrappers/zero_dim.jl
@@ -0,0 +1,56 @@
+# This file is a part of RobustNeuralNetworks.jl. License is MIT: https://github.com/acfr/RobustNeuralNetworks.jl/blob/main/LICENSE 
+
+using Flux
+using Random
+using RobustNeuralNetworks
+using Test
+
+# include("../test_utils.jl")
+
+function (m::AbstractREN{T})(
+    xt::AbstractVecOrMat, 
+    ut::AbstractVecOrMat,
+    explicit::ExplicitRENParams{T}
+) where T
+
+    println("testing")
+    b = explicit.C1 * xt + explicit.D12 * ut .+ explicit.bv
+    wt = RobustNeuralNetworks.tril_eq_layer(m.nl, explicit.D11, b)
+    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ explicit.bx
+    yt = explicit.C2 * xt + explicit.D21 * wt + explicit.D22 * ut .+ explicit.by
+
+    return xt1, yt
+end
+
+"""
+Test that backpropagation runs and parameters change
+"""
+batches = 10
+nu, nx, nv, ny, γ = 4, 5, 0, 2, 10
+ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
+model = DiffREN(ren_ps)
+
+# Dummy data
+us = randn(nu, batches)
+ys = randn(ny, batches)
+data = [(us, ys)]
+
+# Dummy loss function just for testing
+function loss(m, u, y)
+    x0 = init_states(m, size(u,2))
+    x1, y1 = m(x0, u)
+    return Flux.mse(y1, y) + sum(x1.^2)
+end
+
+# Debug batch updates
+opt_state = Flux.setup(Adam(0.01), model)
+gs = Flux.gradient(loss, model, us, ys)
+println()
+
+# # Check if parameters change after a Flux update
+# ps1 = deepcopy(Flux.params(model))
+# opt_state = Flux.setup(Adam(0.01), model)
+# Flux.train!(loss, model, data, opt_state)
+# ps2 = Flux.params(model)
+
+# @test !any(ps1 .≈ ps2)

From a92ea1e8c4511bdb32e3fdfbd3fa3262d34f5b63 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Thu, 13 Jul 2023 09:07:09 +1000
Subject: [PATCH 6/8] Found a slow fix with if statements on bias vectors

---
 test/Wrappers/zero_dim.jl | 41 +++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/test/Wrappers/zero_dim.jl b/test/Wrappers/zero_dim.jl
index 1a2d4d0..3f39f02 100644
--- a/test/Wrappers/zero_dim.jl
+++ b/test/Wrappers/zero_dim.jl
@@ -5,18 +5,19 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-# include("../test_utils.jl")
-
 function (m::AbstractREN{T})(
     xt::AbstractVecOrMat, 
     ut::AbstractVecOrMat,
     explicit::ExplicitRENParams{T}
 ) where T
 
-    println("testing")
-    b = explicit.C1 * xt + explicit.D12 * ut .+ explicit.bv
+    # Allocate bias vectors to avoid error when nv = 0 or nx = 0
+    bv = (m.nv == 0) ? 0 : explicit.bv
+    bx = (m.nx == 0) ? 0 : explicit.bx
+
+    b = explicit.C1 * xt + explicit.D12 * ut .+ bv
     wt = RobustNeuralNetworks.tril_eq_layer(m.nl, explicit.D11, b)
-    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ explicit.bx
+    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ bx
     yt = explicit.C2 * xt + explicit.D21 * wt + explicit.D22 * ut .+ explicit.by
 
     return xt1, yt
@@ -26,9 +27,9 @@ end
 Test that backpropagation runs and parameters change
 """
 batches = 10
-nu, nx, nv, ny, γ = 4, 5, 0, 2, 10
-ren_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
-model = DiffREN(ren_ps)
+nu, nx, nv, ny = 4, 0, 0, 2
+γ = 10
+model_ps = LipschitzRENParams{Float64}(nu, nx, nv, ny, γ)
 
 # Dummy data
 us = randn(nu, batches)
@@ -36,21 +37,23 @@ ys = randn(ny, batches)
 data = [(us, ys)]
 
 # Dummy loss function just for testing
-function loss(m, u, y)
+function loss(model_ps, u, y)
+    m = REN(model_ps)
     x0 = init_states(m, size(u,2))
     x1, y1 = m(x0, u)
     return Flux.mse(y1, y) + sum(x1.^2)
 end
 
-# Debug batch updates
-opt_state = Flux.setup(Adam(0.01), model)
-gs = Flux.gradient(loss, model, us, ys)
-println()
+# Test it
+model = REN(model_ps);
+x0 = init_states(model, batches)
 
-# # Check if parameters change after a Flux update
-# ps1 = deepcopy(Flux.params(model))
-# opt_state = Flux.setup(Adam(0.01), model)
-# Flux.train!(loss, model, data, opt_state)
-# ps2 = Flux.params(model)
+# @btime model(x0, us)
+# @btime Flux.gradient(loss, model_ps, us, ys)
+# println()
 
-# @test !any(ps1 .≈ ps2)
+# Make sure batch update actually runs
+opt_state = Flux.setup(Adam(0.01), model)
+gs = Flux.gradient(loss, model, us, ys)
+Flux.update!(opt_state, model, gs[1])
+@test !isempty(gs)

From ac9099ac99eee80c7879de16cc52f9fda0f67d07 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Thu, 13 Jul 2023 09:34:24 +1000
Subject: [PATCH 7/8] Added check for zero-dim bias vectors

---
 src/Wrappers/REN/ren.jl   |  9 +++++++--
 test/Wrappers/zero_dim.jl | 34 ++++------------------------------
 2 files changed, 11 insertions(+), 32 deletions(-)

diff --git a/src/Wrappers/REN/ren.jl b/src/Wrappers/REN/ren.jl
index 0c1a552..e727f34 100644
--- a/src/Wrappers/REN/ren.jl
+++ b/src/Wrappers/REN/ren.jl
@@ -77,9 +77,14 @@ function (m::AbstractREN{T})(
     explicit::ExplicitRENParams{T}
 ) where T
 
-    b = explicit.C1 * xt + explicit.D12 * ut .+ explicit.bv
+    # Allocate bias vectors to avoid error when nv = 0 or nx = 0
+    # TODO: if statement (or equivalent) makes backpropagation slower. Can we avoid this?
+    bv = (m.nv == 0) ? 0 : explicit.bv
+    bx = (m.nx == 0) ? 0 : explicit.bx
+
+    b = explicit.C1 * xt + explicit.D12 * ut .+ bv
     wt = tril_eq_layer(m.nl, explicit.D11, b)
-    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ explicit.bx
+    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ bx
     yt = explicit.C2 * xt + explicit.D21 * wt + explicit.D22 * ut .+ explicit.by
 
     return xt1, yt
diff --git a/test/Wrappers/zero_dim.jl b/test/Wrappers/zero_dim.jl
index 3f39f02..8591ffa 100644
--- a/test/Wrappers/zero_dim.jl
+++ b/test/Wrappers/zero_dim.jl
@@ -5,26 +5,8 @@ using Random
 using RobustNeuralNetworks
 using Test
 
-function (m::AbstractREN{T})(
-    xt::AbstractVecOrMat, 
-    ut::AbstractVecOrMat,
-    explicit::ExplicitRENParams{T}
-) where T
-
-    # Allocate bias vectors to avoid error when nv = 0 or nx = 0
-    bv = (m.nv == 0) ? 0 : explicit.bv
-    bx = (m.nx == 0) ? 0 : explicit.bx
-
-    b = explicit.C1 * xt + explicit.D12 * ut .+ bv
-    wt = RobustNeuralNetworks.tril_eq_layer(m.nl, explicit.D11, b)
-    xt1 = explicit.A * xt + explicit.B1 * wt + explicit.B2 * ut .+ bx
-    yt = explicit.C2 * xt + explicit.D21 * wt + explicit.D22 * ut .+ explicit.by
-
-    return xt1, yt
-end
-
 """
-Test that backpropagation runs and parameters change
+Test that backpropagation runs when nx = 0 and nv = 0
 """
 batches = 10
 nu, nx, nv, ny = 4, 0, 0, 2
@@ -44,16 +26,8 @@ function loss(model_ps, u, y)
     return Flux.mse(y1, y) + sum(x1.^2)
 end
 
-# Test it
-model = REN(model_ps);
-x0 = init_states(model, batches)
-
-# @btime model(x0, us)
-# @btime Flux.gradient(loss, model_ps, us, ys)
-# println()
-
 # Make sure batch update actually runs
-opt_state = Flux.setup(Adam(0.01), model)
-gs = Flux.gradient(loss, model, us, ys)
-Flux.update!(opt_state, model, gs[1])
+opt_state = Flux.setup(Adam(0.01), model_ps)
+gs = Flux.gradient(loss, model_ps, us, ys)
+Flux.update!(opt_state, model_ps, gs[1])
 @test !isempty(gs)

From be6281e9dd88b0944eb0ce2822d98fefeacb5015 Mon Sep 17 00:00:00 2001
From: nicBarbara <nbar5346@uni.sydney.edu.au>
Date: Thu, 13 Jul 2023 09:49:45 +1000
Subject: [PATCH 8/8] Added test for nx = 0, nv = 0

---
 src/RobustNeuralNetworks.jl | 2 +-
 test/runtests.jl            | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/RobustNeuralNetworks.jl b/src/RobustNeuralNetworks.jl
index 40a77d9..fbe457b 100644
--- a/src/RobustNeuralNetworks.jl
+++ b/src/RobustNeuralNetworks.jl
@@ -6,7 +6,7 @@ module RobustNeuralNetworks
 
 using Flux
 using LinearAlgebra
-using MatrixEquations: lyapd, plyapd
+using MatrixEquations: lyapd
 using Random
 using Zygote: pullback, Buffer
 using Zygote: @adjoint
diff --git a/test/runtests.jl b/test/runtests.jl
index 2651dab..0c9ec15 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -20,5 +20,6 @@ using Test
     include("Wrappers/wrap_ren.jl")
     include("Wrappers/diff_ren.jl")
     include("Wrappers/diff_lbdn.jl")
+    include("Wrappers/zero_dim.jl")
 
 end