From fe207a5ca5a3fda9de8ce2faa65913214c0a2be2 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Sat, 27 Sep 2025 18:35:08 +0200
Subject: [PATCH 01/10] Implement hcat evaluation WIP

Notes: * Gradient is not correct yet (test fails)
* hcat cannot be evaluated as it is, since the final output must be scalar and we do not have any operation to reduce the matrix to scalar at the moment
---
 src/reverse_mode.jl | 53 +++++++++++++++++++++++++++++++++++++++++++++
 src/sizes.jl        | 10 +++++++++
 test/ArrayDiff.jl   | 32 ++++++++++++++++++++++++++-
 3 files changed, 94 insertions(+), 1 deletion(-)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index baab908..268ecd1 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -247,6 +247,24 @@ function _forward_eval(
                     tmp_dot += v1 * v2
                 end
                 @s f.forward_storage[k] = tmp_dot
+            elseif node.index == 12 # hcat
+                idx1, idx2 = children_indices
+                ix1 = children_arr[idx1]
+                ix2 = children_arr[idx2]
+                nb_cols1 = f.sizes.ndims[ix1] <= 1 ? 1 : _size(f.sizes, ix1, 2)
+                col_size = _size(f.sizes, k, 1)
+                for j in _eachindex(f.sizes, k)
+                    col = (j - 1) ÷ col_size + 1
+                    if col <= nb_cols1
+                        @j f.partials_storage[ix1] = one(T)
+                        val = @j f.forward_storage[ix1]
+                        @j f.forward_storage[k] = val
+                    else
+                        @j f.partials_storage[ix2] = one(T)
+                        val = @j f.forward_storage[ix2]
+                        @j f.forward_storage[k] = val
+                    end
+                end
             else # atan, min, max
                 f_input = _UnsafeVectorView(d.jac_storage, N)
                 ∇f = _UnsafeVectorView(d.user_output_buffer, N)
@@ -323,6 +341,18 @@ function _forward_eval(
             f.partials_storage[rhs] = zero(T)
         end
     end
+    # This function is written assuming that the final output is scalar.
+    # Therefore cannot return the matrix, so I guess I return it's first entry only, 
+    # as long as sum or matx-vect products are not implemented.
+
+    #println("Last node ", f.nodes[1].index)
+    #if f.nodes[1].index == 12
+    #    mtx = reshape(
+    #        f.forward_storage[_storage_range(f.sizes, 1)],
+    #        f.sizes.size[1:f.sizes.ndims[1]]...,
+    #    )
+    #    return mtx
+    #end
     return f.forward_storage[1]
 end
 
@@ -379,6 +409,29 @@ function _reverse_eval(f::_SubexpressionStorage)
                         end
                     end
                     continue
+                elseif op == :hcat
+                    total_cols = 0
+                    for c_idx in children_indices
+                        total_cols += f.sizes.ndims[children_arr[c_idx]] <= 1 ?
+                            1 : _size(f.sizes, children_arr[c_idx], 2)
+                    end
+                    row_size = _size(f.sizes, k, 1)
+                    col_start = 1
+                    for c_idx in children_indices
+                        child = children_arr[c_idx]
+                        child_col_size = f.sizes.ndims[child] <= 1 ?
+                            1 : _size(f.sizes, child, 2)
+                        for i in 1:row_size
+                            for j in 1:child_col_size
+                                rev_parent_j =
+                                    @j f.reverse_storage[k]
+                                # partial is 1 so we can ignore it
+                                @j f.reverse_storage[child] = rev_parent_j
+                            end
+                        end
+                        col_start += child_col_size
+                    end
+                    continue
                 end
             end
         elseif node.type != MOI.Nonlinear.NODE_CALL_UNIVARIATE
diff --git a/src/sizes.jl b/src/sizes.jl
index f5ca9dd..f5bd67a 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -184,6 +184,16 @@ function _infer_sizes(
             elseif op == :+ || op == :-
                 # TODO assert all arguments have same size
                 _copy_size!(sizes, k, children_arr[first(children_indices)])
+            elseif op == :hcat
+                total_cols = 0
+                for c_idx in children_indices
+                    total_cols += sizes.ndims[children_arr[c_idx]] <= 1 ?
+                        1 : _size(sizes, children_arr[c_idx], 2)
+                end
+                child_shape = _size(sizes, children_arr[first(children_indices)])
+                shape = sizes.ndims[children_arr[first(children_indices)]] <= 2 ?
+                    (child_shape[1], total_cols) : (child_shape[1], total_cols, child_shape[3:end]...)
+                _add_size!(sizes, k, tuple(shape...))
             elseif op == :*
                 # TODO assert compatible sizes and all ndims should be 0 or 2
                 first_matrix = findfirst(children_indices) do i
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 6dd1f6a..7fedbb6 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -64,6 +64,36 @@ function test_objective_dot_bivariate()
     return
 end
 
+function test_objective_matrix()
+    model = Nonlinear.Model()
+    x1 = MOI.VariableIndex(1)
+    x2 = MOI.VariableIndex(2)
+    x3 = MOI.VariableIndex(3)
+    x4 = MOI.VariableIndex(4)
+    Nonlinear.set_objective(
+        model,
+        :(hcat([$x1, $x3], [$x2, $x4]))
+    )
+    evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
+    MOI.initialize(evaluator, [:Grad])
+    sizes = evaluator.backend.objective.expr.sizes
+    @test sizes.ndims == [2, 1, 0, 0, 1, 0, 0]
+    @test sizes.size_offset == [2, 1, 0, 0, 0, 0, 0]
+    @test sizes.size == [2, 2, 2, 2]
+    @test sizes.storage_offset == [0, 4, 6, 7, 8, 10, 11, 12]
+    x1 = 1.0
+    x2 = 2.0
+    x3 = 3.0
+    x4 = 4.0
+    println(MOI.eval_objective(evaluator, [x1, x2, x3, x4]))
+    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) ==
+          1.0
+    g = ones(4)
+    MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
+    @test g == [1.0, 0.0, 0.0, 0.0]
+    return
+end
+
 end  # module
 
-TestArrayDiff.runtests()
+TestArrayDiff.runtests()
\ No newline at end of file

From b59f8e6a80dfb9dc00d23ff208b0329e3e1ad196 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Sun, 28 Sep 2025 16:23:05 +0200
Subject: [PATCH 02/10] Fix hcat eval and gradient

0dim still to be fixed (hcat scalars into a row)
---
 src/reverse_mode.jl | 62 ++++++++++++++++++++++-----------------------
 test/ArrayDiff.jl   | 16 ++++++------
 2 files changed, 39 insertions(+), 39 deletions(-)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 268ecd1..632a9db 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -252,18 +252,16 @@ function _forward_eval(
                 ix1 = children_arr[idx1]
                 ix2 = children_arr[idx2]
                 nb_cols1 = f.sizes.ndims[ix1] <= 1 ? 1 : _size(f.sizes, ix1, 2)
-                col_size = _size(f.sizes, k, 1)
-                for j in _eachindex(f.sizes, k)
-                    col = (j - 1) ÷ col_size + 1
-                    if col <= nb_cols1
-                        @j f.partials_storage[ix1] = one(T)
-                        val = @j f.forward_storage[ix1]
-                        @j f.forward_storage[k] = val
-                    else
-                        @j f.partials_storage[ix2] = one(T)
-                        val = @j f.forward_storage[ix2]
-                        @j f.forward_storage[k] = val
-                    end
+                col_size = f.sizes.ndims[ix1] == 0 ? 1 : _size(f.sizes, k, 1)
+                for j in _eachindex(f.sizes, ix1)
+                    @j f.partials_storage[ix1] = one(T)
+                    val = @j f.forward_storage[ix1]
+                    @j f.forward_storage[k] = val
+                end
+                for j in _eachindex(f.sizes, ix2)
+                    @j f.partials_storage[ix2] = one(T)
+                    val = @j f.forward_storage[ix2]
+                    _setindex!(f.forward_storage, val, f.sizes, k, j + nb_cols1 * col_size)
                 end
             else # atan, min, max
                 f_input = _UnsafeVectorView(d.jac_storage, N)
@@ -410,26 +408,28 @@ function _reverse_eval(f::_SubexpressionStorage)
                     end
                     continue
                 elseif op == :hcat
-                    total_cols = 0
-                    for c_idx in children_indices
-                        total_cols += f.sizes.ndims[children_arr[c_idx]] <= 1 ?
-                            1 : _size(f.sizes, children_arr[c_idx], 2)
+                    idx1, idx2 = children_indices
+                    ix1 = children_arr[idx1]
+                    ix2 = children_arr[idx2]
+                    nb_cols1 = f.sizes.ndims[ix1] <= 1 ? 1 : _size(f.sizes, ix1, 2)
+                    col_size = f.sizes.ndims[ix1] == 0 ? 1 : _size(f.sizes, k, 1)
+                    for j in _eachindex(f.sizes, ix1)
+                        partial = @j f.partials_storage[ix1]
+                        val = ifelse(
+                            _getindex(f.reverse_storage, f.sizes, k, j) == 0.0 && !isfinite(partial),
+                            _getindex(f.reverse_storage, f.sizes, k, j),
+                            _getindex(f.reverse_storage, f.sizes, k, j) * partial,
+                        )
+                        @j f.reverse_storage[ix1] = val
                     end
-                    row_size = _size(f.sizes, k, 1)
-                    col_start = 1
-                    for c_idx in children_indices
-                        child = children_arr[c_idx]
-                        child_col_size = f.sizes.ndims[child] <= 1 ?
-                            1 : _size(f.sizes, child, 2)
-                        for i in 1:row_size
-                            for j in 1:child_col_size
-                                rev_parent_j =
-                                    @j f.reverse_storage[k]
-                                # partial is 1 so we can ignore it
-                                @j f.reverse_storage[child] = rev_parent_j
-                            end
-                        end
-                        col_start += child_col_size
+                    for j in _eachindex(f.sizes, ix2)
+                        partial = @j f.partials_storage[ix2]
+                        val = ifelse(
+                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size) == 0.0 && !isfinite(partial),
+                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size),
+                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size) * partial,
+                        )
+                        @j f.reverse_storage[ix2] = val
                     end
                     continue
                 end
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 7fedbb6..6b16ea9 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -64,7 +64,7 @@ function test_objective_dot_bivariate()
     return
 end
 
-function test_objective_matrix()
+function test_objective_hcat_1dim()
     model = Nonlinear.Model()
     x1 = MOI.VariableIndex(1)
     x2 = MOI.VariableIndex(2)
@@ -72,25 +72,25 @@ function test_objective_matrix()
     x4 = MOI.VariableIndex(4)
     Nonlinear.set_objective(
         model,
-        :(hcat([$x1, $x3], [$x2, $x4]))
+        :(dot(hcat([$x1], [$x3]), hcat([$x2], [$x4])))
     )
     evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
     MOI.initialize(evaluator, [:Grad])
     sizes = evaluator.backend.objective.expr.sizes
-    @test sizes.ndims == [2, 1, 0, 0, 1, 0, 0]
-    @test sizes.size_offset == [2, 1, 0, 0, 0, 0, 0]
-    @test sizes.size == [2, 2, 2, 2]
-    @test sizes.storage_offset == [0, 4, 6, 7, 8, 10, 11, 12]
+    @test sizes.ndims == [0, 2, 1, 0, 1, 0, 2, 1, 0, 1, 0]
+    @test sizes.size_offset == [0, 6, 5, 0, 4, 0, 2, 1, 0, 0, 0]
+    @test sizes.size == [1, 1, 1, 2, 1, 1, 1, 2]
+    @test sizes.storage_offset == [0, 1, 3, 4, 5, 6, 7, 9, 10, 11, 12, 13]
     x1 = 1.0
     x2 = 2.0
     x3 = 3.0
     x4 = 4.0
     println(MOI.eval_objective(evaluator, [x1, x2, x3, x4]))
     @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) ==
-          1.0
+          14.0
     g = ones(4)
     MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
-    @test g == [1.0, 0.0, 0.0, 0.0]
+    @test g == [2.0, 1.0, 4.0, 3.0]
     return
 end
 

From 0e90752fe1e67e66af6104fa4f2a7792ab940624 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Sun, 28 Sep 2025 18:19:21 +0200
Subject: [PATCH 03/10] Fix 0dim

Meaning a row vector can now be written with [1 2] or hcat(1, 2)
---
 src/sizes.jl      | 11 ++++++++---
 test/ArrayDiff.jl | 30 ++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/sizes.jl b/src/sizes.jl
index f5bd67a..4a44332 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -190,9 +190,14 @@ function _infer_sizes(
                     total_cols += sizes.ndims[children_arr[c_idx]] <= 1 ?
                         1 : _size(sizes, children_arr[c_idx], 2)
                 end
-                child_shape = _size(sizes, children_arr[first(children_indices)])
-                shape = sizes.ndims[children_arr[first(children_indices)]] <= 2 ?
-                    (child_shape[1], total_cols) : (child_shape[1], total_cols, child_shape[3:end]...)
+                if sizes.ndims[children_arr[first(children_indices)]] == 0
+                    shape = (1, total_cols)
+                elseif sizes.ndims[children_arr[first(children_indices)]] <= 2
+                    shape = ( _size(sizes, children_arr[first(children_indices)], 1), total_cols)
+                else 
+                    child_shape = _size(sizes, children_arr[first(children_indices)])
+                    shape = (child_shape[1], total_cols, child_shape[3:end]...)
+                end
                 _add_size!(sizes, k, tuple(shape...))
             elseif op == :*
                 # TODO assert compatible sizes and all ndims should be 0 or 2
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 6b16ea9..c388443 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -64,6 +64,36 @@ function test_objective_dot_bivariate()
     return
 end
 
+function test_objective_hcat_0dim()
+    model = Nonlinear.Model()
+    x1 = MOI.VariableIndex(1)
+    x2 = MOI.VariableIndex(2)
+    x3 = MOI.VariableIndex(3)
+    x4 = MOI.VariableIndex(4)
+    Nonlinear.set_objective(
+        model,
+        :(dot([$x1 $x3], [$x2 $x4]))
+    )
+    evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
+    MOI.initialize(evaluator, [:Grad])
+    sizes = evaluator.backend.objective.expr.sizes
+    @test sizes.ndims == [0, 2, 0, 0, 2, 0, 0]
+    @test sizes.size_offset == [0, 2, 0, 0, 0, 0, 0]
+    @test sizes.size == [1, 2, 1, 2]
+    @test sizes.storage_offset == [0, 1, 3, 4, 5, 7, 8, 9]
+    x1 = 1.0
+    x2 = 2.0
+    x3 = 3.0
+    x4 = 4.0
+    println(MOI.eval_objective(evaluator, [x1, x2, x3, x4]))
+    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) ==
+          14.0
+    g = ones(4)
+    MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
+    @test g == [2.0, 1.0, 4.0, 3.0]
+    return
+end
+
 function test_objective_hcat_1dim()
     model = Nonlinear.Model()
     x1 = MOI.VariableIndex(1)

From 7842545a2e933448edf31d8ca1845d8f78569206 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Mon, 29 Sep 2025 09:20:32 +0200
Subject: [PATCH 04/10] Add Revise to toml (temporary)

---
 test/Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/test/Project.toml b/test/Project.toml
index 7ed96af..0e679fe 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -2,5 +2,6 @@
 ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

From 7576b515beaaf3db573b7cc504778bf3b817b0f8 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Mon, 29 Sep 2025 09:42:10 +0200
Subject: [PATCH 05/10] Fix quick pb I introduced when merging

---
 src/reverse_mode.jl | 2 ++
 test/ArrayDiff.jl   | 4 +++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 2238614..c12e4be 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -262,6 +262,7 @@ function _forward_eval(
                     @j f.partials_storage[ix2] = one(T)
                     val = @j f.forward_storage[ix2]
                     _setindex!(f.forward_storage, val, f.sizes, k, j + nb_cols1 * col_size)
+                end
             elseif node.index == 14 # norm 
                 ix = children_arr[children_indices[1]]
                 tmp_norm_squared = zero(T)
@@ -445,6 +446,7 @@ function _reverse_eval(f::_SubexpressionStorage)
                             _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size) * partial,
                         )
                         @j f.reverse_storage[ix2] = val
+                    end
                 elseif op == :norm
                     # Node `k` is scalar, the jacobian w.r.t. the vectorized input
                     # child is a row vector whose entries are stored in `f.partials_storage`
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 3b4776d..093f0c9 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -121,6 +121,8 @@ function test_objective_hcat_1dim()
     g = ones(4)
     MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
     @test g == [2.0, 1.0, 4.0, 3.0]
+    return
+end
   
 function test_objective_norm_univariate()
     model = Nonlinear.Model()
@@ -168,4 +170,4 @@ end
 
 end  # module
 
-TestArrayDiff.runtests()
\ No newline at end of file
+TestArrayDiff.runtests() 
\ No newline at end of file

From 5fe52642dbf1a0b9d1ab91da1c66f52650d2871b Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Mon, 29 Sep 2025 10:00:37 +0200
Subject: [PATCH 06/10] Add missing continue

To prevent checking if hcat has same size as children
---
 src/reverse_mode.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index c12e4be..8c53731 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -447,6 +447,7 @@ function _reverse_eval(f::_SubexpressionStorage)
                         )
                         @j f.reverse_storage[ix2] = val
                     end
+                    continue
                 elseif op == :norm
                     # Node `k` is scalar, the jacobian w.r.t. the vectorized input
                     # child is a row vector whose entries are stored in `f.partials_storage`

From bc3c92483d57f458631698862332c93f15e9f812 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Wed, 1 Oct 2025 11:21:19 +0200
Subject: [PATCH 07/10] Update test/ArrayDiff.jl

---
 test/ArrayDiff.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 093f0c9..0b8529a 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -170,4 +170,4 @@ end
 
 end  # module
 
-TestArrayDiff.runtests() 
\ No newline at end of file
+TestArrayDiff.runtests()
\ No newline at end of file

From b4770da6faf0a296574d45ffe25392b69efecc6d Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Thu, 2 Oct 2025 07:07:10 +0200
Subject: [PATCH 08/10] Fix format

---
 src/reverse_mode.jl | 43 ++++++++++++++++++++++++++++++++++---------
 src/sizes.jl        | 15 ++++++++++-----
 test/ArrayDiff.jl   | 17 ++++++-----------
 3 files changed, 50 insertions(+), 25 deletions(-)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 8c53731..6cdc70f 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -261,7 +261,13 @@ function _forward_eval(
                 for j in _eachindex(f.sizes, ix2)
                     @j f.partials_storage[ix2] = one(T)
                     val = @j f.forward_storage[ix2]
-                    _setindex!(f.forward_storage, val, f.sizes, k, j + nb_cols1 * col_size)
+                    _setindex!(
+                        f.forward_storage,
+                        val,
+                        f.sizes,
+                        k,
+                        j + nb_cols1 * col_size,
+                    )
                 end
             elseif node.index == 14 # norm 
                 ix = children_arr[children_indices[1]]
@@ -427,23 +433,42 @@ function _reverse_eval(f::_SubexpressionStorage)
                     idx1, idx2 = children_indices
                     ix1 = children_arr[idx1]
                     ix2 = children_arr[idx2]
-                    nb_cols1 = f.sizes.ndims[ix1] <= 1 ? 1 : _size(f.sizes, ix1, 2)
-                    col_size = f.sizes.ndims[ix1] == 0 ? 1 : _size(f.sizes, k, 1)
+                    nb_cols1 =
+                        f.sizes.ndims[ix1] <= 1 ? 1 : _size(f.sizes, ix1, 2)
+                    col_size =
+                        f.sizes.ndims[ix1] == 0 ? 1 : _size(f.sizes, k, 1)
                     for j in _eachindex(f.sizes, ix1)
                         partial = @j f.partials_storage[ix1]
                         val = ifelse(
-                            _getindex(f.reverse_storage, f.sizes, k, j) == 0.0 && !isfinite(partial),
+                            _getindex(f.reverse_storage, f.sizes, k, j) ==
+                            0.0 && !isfinite(partial),
                             _getindex(f.reverse_storage, f.sizes, k, j),
-                            _getindex(f.reverse_storage, f.sizes, k, j) * partial,
+                            _getindex(f.reverse_storage, f.sizes, k, j) *
+                            partial,
                         )
                         @j f.reverse_storage[ix1] = val
                     end
                     for j in _eachindex(f.sizes, ix2)
                         partial = @j f.partials_storage[ix2]
                         val = ifelse(
-                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size) == 0.0 && !isfinite(partial),
-                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size),
-                            _getindex(f.reverse_storage, f.sizes, k, j + nb_cols1 * col_size) * partial,
+                            _getindex(
+                                f.reverse_storage,
+                                f.sizes,
+                                k,
+                                j + nb_cols1 * col_size,
+                            ) == 0.0 && !isfinite(partial),
+                            _getindex(
+                                f.reverse_storage,
+                                f.sizes,
+                                k,
+                                j + nb_cols1 * col_size,
+                            ),
+                            _getindex(
+                                f.reverse_storage,
+                                f.sizes,
+                                k,
+                                j + nb_cols1 * col_size,
+                            ) * partial,
                         )
                         @j f.reverse_storage[ix2] = val
                     end
@@ -461,7 +486,7 @@ function _reverse_eval(f::_SubexpressionStorage)
                             rev_parent,
                             rev_parent * partial,
                         )
-                        @j f.reverse_storage[ix] = val  
+                        @j f.reverse_storage[ix] = val
                     end
                     continue
                 end
diff --git a/src/sizes.jl b/src/sizes.jl
index 518b040..002d290 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -189,15 +189,20 @@ function _infer_sizes(
             elseif op == :hcat
                 total_cols = 0
                 for c_idx in children_indices
-                    total_cols += sizes.ndims[children_arr[c_idx]] <= 1 ?
-                        1 : _size(sizes, children_arr[c_idx], 2)
+                    total_cols +=
+                        sizes.ndims[children_arr[c_idx]] <= 1 ? 1 :
+                        _size(sizes, children_arr[c_idx], 2)
                 end
                 if sizes.ndims[children_arr[first(children_indices)]] == 0
                     shape = (1, total_cols)
                 elseif sizes.ndims[children_arr[first(children_indices)]] <= 2
-                    shape = ( _size(sizes, children_arr[first(children_indices)], 1), total_cols)
-                else 
-                    child_shape = _size(sizes, children_arr[first(children_indices)])
+                    shape = (
+                        _size(sizes, children_arr[first(children_indices)], 1),
+                        total_cols,
+                    )
+                else
+                    child_shape =
+                        _size(sizes, children_arr[first(children_indices)])
                     shape = (child_shape[1], total_cols, child_shape[3:end]...)
                 end
                 _add_size!(sizes, k, tuple(shape...))
diff --git a/test/ArrayDiff.jl b/test/ArrayDiff.jl
index 0b8529a..4bd270b 100644
--- a/test/ArrayDiff.jl
+++ b/test/ArrayDiff.jl
@@ -70,10 +70,7 @@ function test_objective_hcat_0dim()
     x2 = MOI.VariableIndex(2)
     x3 = MOI.VariableIndex(3)
     x4 = MOI.VariableIndex(4)
-    Nonlinear.set_objective(
-        model,
-        :(dot([$x1 $x3], [$x2 $x4]))
-    )
+    Nonlinear.set_objective(model, :(dot([$x1 $x3], [$x2 $x4])))
     evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
     MOI.initialize(evaluator, [:Grad])
     sizes = evaluator.backend.objective.expr.sizes
@@ -86,8 +83,7 @@ function test_objective_hcat_0dim()
     x3 = 3.0
     x4 = 4.0
     println(MOI.eval_objective(evaluator, [x1, x2, x3, x4]))
-    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) ==
-          14.0
+    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) == 14.0
     g = ones(4)
     MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
     @test g == [2.0, 1.0, 4.0, 3.0]
@@ -102,7 +98,7 @@ function test_objective_hcat_1dim()
     x4 = MOI.VariableIndex(4)
     Nonlinear.set_objective(
         model,
-        :(dot(hcat([$x1], [$x3]), hcat([$x2], [$x4])))
+        :(dot(hcat([$x1], [$x3]), hcat([$x2], [$x4]))),
     )
     evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
     MOI.initialize(evaluator, [:Grad])
@@ -116,14 +112,13 @@ function test_objective_hcat_1dim()
     x3 = 3.0
     x4 = 4.0
     println(MOI.eval_objective(evaluator, [x1, x2, x3, x4]))
-    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) ==
-          14.0
+    @test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) == 14.0
     g = ones(4)
     MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
     @test g == [2.0, 1.0, 4.0, 3.0]
     return
 end
-  
+
 function test_objective_norm_univariate()
     model = Nonlinear.Model()
     x = MOI.VariableIndex(1)
@@ -170,4 +165,4 @@ end
 
 end  # module
 
-TestArrayDiff.runtests()
\ No newline at end of file
+TestArrayDiff.runtests()

From 8e32bf01eaa5a9d1b0fbc3ba1246b6c2f51ecf12 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Thu, 2 Oct 2025 09:04:00 +0200
Subject: [PATCH 09/10] Remove Revise and simplify hcat

* Remove Revise from toml as mentionned in review
* Remove hcat case 'more than 2 dims', as we do not pretend to handle it right now
---
 src/sizes.jl      | 9 ++++-----
 test/Project.toml | 1 -
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/sizes.jl b/src/sizes.jl
index 002d290..47aefbd 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -195,15 +195,14 @@ function _infer_sizes(
                 end
                 if sizes.ndims[children_arr[first(children_indices)]] == 0
                     shape = (1, total_cols)
-                elseif sizes.ndims[children_arr[first(children_indices)]] <= 2
+                else
+                    @assert sizes.ndims[children_arr[first(
+                        children_indices,
+                    )]] <= 2 "Hcat with ndims > 2 is not supported yet"
                     shape = (
                         _size(sizes, children_arr[first(children_indices)], 1),
                         total_cols,
                     )
-                else
-                    child_shape =
-                        _size(sizes, children_arr[first(children_indices)])
-                    shape = (child_shape[1], total_cols, child_shape[3:end]...)
                 end
                 _add_size!(sizes, k, tuple(shape...))
             elseif op == :*
diff --git a/test/Project.toml b/test/Project.toml
index 0e679fe..7ed96af 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -2,6 +2,5 @@
 ArrayDiff = "c45fa1ca-6901-44ac-ae5b-5513a4852d50"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
-Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

From 8983f4df7051509911408d5d97fa8f28aa987308 Mon Sep 17 00:00:00 2001
From: Sophie L <sophie.lequeu@student.uclouvain.be>
Date: Thu, 2 Oct 2025 10:04:33 +0200
Subject: [PATCH 10/10] Delete personal comments

---
 src/reverse_mode.jl | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 6cdc70f..16ef670 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -361,18 +361,6 @@ function _forward_eval(
             f.partials_storage[rhs] = zero(T)
         end
     end
-    # This function is written assuming that the final output is scalar.
-    # Therefore cannot return the matrix, so I guess I return it's first entry only, 
-    # as long as sum or matx-vect products are not implemented.
-
-    #println("Last node ", f.nodes[1].index)
-    #if f.nodes[1].index == 12
-    #    mtx = reshape(
-    #        f.forward_storage[_storage_range(f.sizes, 1)],
-    #        f.sizes.size[1:f.sizes.ndims[1]]...,
-    #    )
-    #    return mtx
-    #end
     return f.forward_storage[1]
 end