diff --git a/.github/workflows/build-foam.yml b/.github/workflows/build-foam.yml
index a0284d057..b7ad85d26 100644
--- a/.github/workflows/build-foam.yml
+++ b/.github/workflows/build-foam.yml
@@ -69,6 +69,12 @@ jobs:
          key: build_PR_${{ github.event.pull_request.number }}_${{matrix.OF.path}}
          path: ${{github.workspace}}/build
 
+    - name: Cache FOAM_USER_LIBBIN
+      uses: actions/cache@v4
+      with:
+         key: FOAM_USER_LIBBIN-${{ matrix.OF.path }}-${{ github.sha }}
+         path: ${{matrix.OF.foam_user_libbin}}
+
     - name: Configure CMake
       # Use a bash shell so we can use the same syntax for environment variable
       # access regardless of the host operating system
@@ -76,29 +82,19 @@ jobs:
       run: |
         source /root/OpenFOAM/${{matrix.OF.path}}/etc/bashrc
         mkdir -p ${{matrix.OF.foam_user_libbin}}
-        mkdir -p build
-        cd build
-        [ -d "/github/home/$GINKGO_CHECKOUT_VERSION" ] && cp -rp /github/home/$GINKGO_CHECKOUT_VERSION third_party
-        cmake -GNinja \
-          -DOGL_ALLOW_REFERENCE_ONLY=On \
-          -DOGL_BUILD_UNITTEST=ON \
-          -DCMAKE_BUILD_TYPE=$BUILD_TYPE \
-          -DCMAKE_INSTALL_PREFIX=${{matrix.OF.foam_user_libbin}} ..
+        cmake --preset develop
 
     - name: Build OGL
-      working-directory: ${{github.workspace}}/build
+      working-directory: ${{github.workspace}}
       run: |
         source /root/OpenFOAM/${{matrix.OF.path}}/etc/bashrc
-        cmake --build . --config $BUILD_TYPE
+        cmake  --build --preset develop
 
     - name: Install OGL
       working-directory: ${{github.workspace}}
       run: |
         source /root/OpenFOAM/${{matrix.OF.path}}/etc/bashrc
-        cmake --install ./build
-        ls $FOAM_USER_LIBBIN
-        ls ${{matrix.OF.foam_user_libbin}}
-
+        cmake --build --preset develop --target install
 
     - name: Check folders
       working-directory: ${{github.workspace}}
diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
index 85defba69..01e1a89d5 100644
--- a/.github/workflows/integration-tests.yml
+++ b/.github/workflows/integration-tests.yml
@@ -89,7 +89,7 @@ jobs:
         export LD_LIBRARY_PATH=$FOAM_USER_LIBBIN:$WM_PROJECT_DIR/platforms/linux64GccDPInt32Opt/lib:$WM_PROJECT_DIR/platforms/linux64GccDPInt32Opt/lib/sys-openmpi/:$WM_PROJECT_DIR/platforms/linux64GccDPInt32Opt/lib/openmpi-system:$WM_PROJECT_DIR/platforms/linux64GccDPInt32Opt/lib/dummy/:$LD_LIBRARY_PATH
         mkdir -p integration && cd integration
         export CYCLIC_CASE=${{inputs.cyclic_case}}
-        obr init -gy --config /__w/OGL/OGL/test/integration.yaml
+        obr init -g --config /__w/OGL/OGL/test/integration.yaml
         cat workspace/*/case/*log
 
   setup_integration_matrix:
@@ -128,7 +128,7 @@ jobs:
     - name: Get Ginkgo checkout version
       shell: bash
       run: |
-        grep -A1 "set(GINKGO_CHECKOUT_VERSION" CMakeLists.txt|tail -n1|grep -o  "[0-9a-z\-]*" > GINKGO_CHECKOUT_VERSION
+        grep -A1 "set(GINKGO_CHECKOUT_VERSION" CMakeLists.txt|tail -n1|grep -o  "[0-9a-z\-_.]*" > GINKGO_CHECKOUT_VERSION
         export GINKGO_CHECKOUT_VERSION=$(cat GINKGO_CHECKOUT_VERSION)
         echo "GINKGO_CHECKOUT_VERSION=$GINKGO_CHECKOUT_VERSION" >> $GITHUB_ENV
         echo "Ginkgo_DIR=${{inputs.foam_user_libbin}}/cmake/Ginkgo" >> $GITHUB_ENV
diff --git a/.github/workflows/static_checks.yaml b/.github/workflows/static_checks.yaml
index df3424044..f5fb4eb28 100644
--- a/.github/workflows/static_checks.yaml
+++ b/.github/workflows/static_checks.yaml
@@ -37,10 +37,10 @@ jobs:
     - name: Create Compilation Database
       run: |
         source /root/OpenFOAM/openfoam/etc/bashrc
-        cmake --preset ninja-cpuonly-release \
+        cmake --preset develop \
           -DCMAKE_CXX_COMPILER=clang++ \
           -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
-        cmake --build --preset ninja-cpuonly-release
+        cmake --build --preset develop
     - uses: actions/upload-artifact@v4
       with:
         name: build-artifact
diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml
index b0faf17b3..6ca823c82 100644
--- a/.github/workflows/unit-test.yml
+++ b/.github/workflows/unit-test.yml
@@ -33,8 +33,9 @@ jobs:
     - uses: actions/checkout@v4
 
     - name: Get Ginkgo checkout version
+      shell: bash
       run: |
-        grep -A1 "set(GINKGO_CHECKOUT_VERSION" CMakeLists.txt|tail -n1|grep -o  "[0-9a-z\-]*" > GINKGO_CHECKOUT_VERSION
+        grep -A1 "set(GINKGO_CHECKOUT_VERSION" CMakeLists.txt|tail -n1|grep -o  "[0-9a-z\-_.]*" > GINKGO_CHECKOUT_VERSION
         export GINKGO_CHECKOUT_VERSION=$(cat GINKGO_CHECKOUT_VERSION)
         echo "GINKGO_CHECKOUT_VERSION=$GINKGO_CHECKOUT_VERSION" >> $GITHUB_ENV
 
@@ -48,13 +49,13 @@ jobs:
     - name: Cache build folder
       uses: actions/cache@v4
       with:
-        key: build-${{ inputs.path }}-${{env.GINKGO_CHECKOUT_VERSION}}
+        # key: build-${{ inputs.path }}-${{env.GINKGO_CHECKOUT_VERSION}}
+        key: build_PR_${{ github.event.pull_request.number }}_${{inputs.path}}
         path: |
           ${{github.workspace}}/build
 
     - name: Test OGL
-      working-directory: ${{github.workspace}}/build
+      working-directory: ${{github.workspace}}
       run: |
         source /root/OpenFOAM/${{inputs.path}}/etc/bashrc
-        ls $FOAM_USER_LIBBIN
-        ./unitTests/matrixConversion
+        ctest --preset develop
diff --git a/CMakePresets.json b/CMakePresets.json
index 401b50010..2659bf43c 100644
--- a/CMakePresets.json
+++ b/CMakePresets.json
@@ -84,6 +84,18 @@
       "configuration": "Develop"
     }
   ],
+  "testPresets": [
+    {
+      "name": "debug",
+      "configurePreset": "debug",
+      "configuration": "Release"
+    },
+    {
+      "name": "develop",
+      "configurePreset": "develop",
+      "configuration": "Develop"
+    }
+  ],
   "workflowPresets": [
     {
       "name": "release",
diff --git a/include/OGL/CommunicationPattern.H b/include/OGL/CommunicationPattern.H
index c3d12ee79..356e0fe18 100644
--- a/include/OGL/CommunicationPattern.H
+++ b/include/OGL/CommunicationPattern.H
@@ -142,7 +142,7 @@ struct CommunicationPattern {
           send_idxs(send_idxs_)
     {
         // every communication ranks needs a set of send_idxs associated
-        ASSERT_EQ(ranks.size(), send_idxs_.size());
+        OGL_ASSERT_EQ(ranks.size(), send_idxs_.size());
 
         for (auto &rows : send_idxs) {
             target_sizes.push_back(rows.size());
diff --git a/include/OGL/DevicePersistent/Base.H b/include/OGL/DevicePersistent/Base.H
index 4c587a1ea..52e42ca58 100644
--- a/include/OGL/DevicePersistent/Base.H
+++ b/include/OGL/DevicePersistent/Base.H
@@ -38,7 +38,7 @@ public:
 
 /* Base class for storing ginkgo and OGL data structures in the Objectregistry
  *
- * This class creates an unitialised Ginkgo array and copies the
+ * This class creates an uninitialised Ginkgo array and copies the
  * shared ptr to the DevicePersistentBase. This will prevent it from
  * being deleted when leaving the scope.
  *
diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index 89066abe5..78e30a81b 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -51,15 +51,14 @@ public:
         return this->dist_mtx_->get_non_local_matrix();
     }
 
-    template <typename InnerType>
     std::shared_ptr<const gko::LinOp> get_local() const
     {
         if (repartitioner_->get_fused()) {
             return dist_mtx_->get_local_matrix();
         } else {
-            return gko::as<InnerType>(gko::as<CombinationMatrix<InnerType>>(
-                                          dist_mtx_->get_local_matrix())
-                                          ->get_operators()[0]);
+            return gko::as<CombinationMatrix<gko::LinOp>>(
+                       dist_mtx_->get_local_matrix())
+                ->get_operators()[0];
         }
     }
 
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 70db234da..682b252d0 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -53,8 +53,8 @@ struct SparsityPattern {
           dim(dim_),
           spans(spans_)
     {
-        ASSERT_EQ(rows.size(), cols.size());
-        ASSERT_EQ(rows.size(), mapping.size());
+        OGL_ASSERT_EQ(rows.size(), cols.size());
+        OGL_ASSERT_EQ(rows.size(), mapping.size());
     }
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec)
@@ -85,16 +85,16 @@ struct SparsityPattern {
           ldu_mapping{exec, static_cast<gko::size_type>(size)},
           dim(dim)
     {
-        ASSERT_EQ(static_cast<size_t>(size), rows.size());
+        OGL_ASSERT_EQ(static_cast<size_t>(size), rows.size());
         std::copy(rows.begin(), rows.end(), row_idxs.get_data());
 
-        ASSERT_EQ(static_cast<size_t>(size), cols.size());
+        OGL_ASSERT_EQ(static_cast<size_t>(size), cols.size());
         std::copy(cols.begin(), cols.end(), col_idxs.get_data());
 
-        ASSERT_EQ(static_cast<size_t>(size), map.size());
+        OGL_ASSERT_EQ(static_cast<size_t>(size), map.size());
         std::copy(map.begin(), map.end(), ldu_mapping.get_data());
 
-        ASSERT_EQ(spans_begin.size(), spans_end.size());
+        OGL_ASSERT_EQ(spans_begin.size(), spans_end.size());
         std::transform(spans_begin.begin(), spans_begin.end(),
                        spans_end.begin(), std::back_inserter(spans),
                        [](const auto begin, const auto end) {
diff --git a/include/OGL/Preconditioner.H b/include/OGL/Preconditioner.H
index dcdf22f3a..3cd838880 100644
--- a/include/OGL/Preconditioner.H
+++ b/include/OGL/Preconditioner.H
@@ -51,7 +51,7 @@ public:
         std::shared_ptr<gko::Executor> device_exec,
         std::unique_ptr<PrecondFactory> precond) const
     {
-        auto local = get_local(gkomatrix);
+        auto local = gko::as<RepartDistMatrix>(gkomatrix)->get_local();
         return gko::share(
             ras::build()
                 .with_generated_local_solver(
diff --git a/include/OGL/common.H b/include/OGL/common.H
index bb1d26265..3e57e3eab 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -90,7 +90,7 @@ namespace Foam {
 
 #define UNUSED(x) (void)(x)
 
-#define ASSERT_EQ(_val1, _val2)                                               \
+#define OGL_ASSERT_EQ(_val1, _val2)                                           \
     if (_val1 != _val2) {                                                     \
         FatalErrorInFunction << "Expected equal values: " << _val1 << " and " \
                              << _val2 << abort(FatalError);                   \
diff --git a/src/CommunicationPattern.C b/src/CommunicationPattern.C
index 63670ad7a..1a4a8a0dc 100644
--- a/src/CommunicationPattern.C
+++ b/src/CommunicationPattern.C
@@ -70,7 +70,7 @@ AllToAllPattern compute_gather_to_owner_counts(
     auto exec = exec_handler.get_device_exec();
     auto comm = *exec_handler.get_communicator().get();
 
-    ASSERT_EQ(total_size, size + padding_before + padding_after);
+    OGL_ASSERT_EQ(total_size, size + padding_before + padding_after);
     label total_ranks{comm.size()};
     label rank{comm.rank()};
     label owner_rank = compute_owner_rank(rank, ranks_per_owner);
@@ -161,26 +161,24 @@ void communicate_values(
     const AllToAllPattern &comm_pattern, const scalar *send_buffer,
     scalar *recv_buffer, bool force_host_buffer, label recv_buffer_size)
 {
-        if (force_host_buffer && src_exec != target_exec) {
-            auto tmp =
-                gko::array<scalar>(src_exec, recv_buffer_size);
+    if (force_host_buffer && src_exec != target_exec) {
+        auto tmp = gko::array<scalar>(src_exec, recv_buffer_size);
 
-            comm->all_to_all_v(src_exec, send_buffer,
-                               comm_pattern.send_counts.data(),
-                               comm_pattern.send_offsets.data(), tmp.get_data(),
-                               comm_pattern.recv_counts.data(),
-                               comm_pattern.recv_offsets.data());
+        comm->all_to_all_v(
+            src_exec, send_buffer, comm_pattern.send_counts.data(),
+            comm_pattern.send_offsets.data(), tmp.get_data(),
+            comm_pattern.recv_counts.data(), comm_pattern.recv_offsets.data());
 
-            auto recv_view = gko::array<scalar>::view(
-                target_exec, recv_buffer_size, recv_buffer);
+        auto recv_view = gko::array<scalar>::view(target_exec, recv_buffer_size,
+                                                  recv_buffer);
 
-            recv_view = tmp;
-        } else {
+        recv_view = tmp;
+    } else {
         comm->all_to_all_v(
             target_exec, send_buffer, comm_pattern.send_counts.data(),
             comm_pattern.send_offsets.data(), recv_buffer,
             comm_pattern.recv_counts.data(), comm_pattern.recv_offsets.data());
-        }
+    }
 }
 
 std::vector<label> gather_labels_to_owner(const ExecutorHandler &exec_handler,
diff --git a/src/DevicePersistent/DistributedMatrixAdapter.C b/src/DevicePersistent/DistributedMatrixAdapter.C
index 3e226228d..340f46af0 100644
--- a/src/DevicePersistent/DistributedMatrixAdapter.C
+++ b/src/DevicePersistent/DistributedMatrixAdapter.C
@@ -10,5 +10,6 @@ namespace Foam {
 
 // to store the std::shared_ptr<T> in the IO registry the type needs to be
 // declared
-defineTemplateTypeNameWithName(DevicePersistentBase<RepartDistMatrix>, "PersistentDistributedMatrix");
+defineTemplateTypeNameWithName(DevicePersistentBase<RepartDistMatrix>,
+                               "PersistentDistributedMatrix");
 }  // namespace Foam
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index 54b3f67a2..f9f6512a6 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -98,10 +98,10 @@ void update_fused_impl(
         dist_A,
     std::shared_ptr<const SparsityPattern> local_sparsity,
     std::shared_ptr<const SparsityPattern> non_local_sparsity,
-    [[maybe_unused]] std::shared_ptr<const CommunicationPattern>
-        src_comm_pattern,
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
     std::vector<InterfaceLocality> local_interfaces)
 {
+    using vec = gko::matrix::Dense<scalar>;
     auto exec = exec_handler.get_ref_exec();
     auto device_exec = exec_handler.get_device_exec();
     auto ranks_per_gpu = repartitioner->get_ranks_per_gpu();
@@ -120,21 +120,20 @@ void update_fused_impl(
         }
     }
 
-
     label tot_local_matrix_nnz = local_matrix_nnz + n_interfaces;
 
     // size + padding has to be local_matrix_nnz
-    // [upper, lower, diag, interfaces] | [upper, lower, diag, interfaces]
+    // [upper, lower, diag] | [upper, lower, diag], [interfaces | interfaces]
     auto diag_comm_pattern = compute_gather_to_owner_counts(
-        exec_handler, ranks_per_gpu, nrows, tot_local_matrix_nnz,
-        local_matrix_nnz - nrows, n_interfaces);
+        exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
+        local_matrix_nnz - nrows, 0);
     label upper_nnz = host_A->get_upper_nnz();
     auto upper_comm_pattern = compute_gather_to_owner_counts(
-        exec_handler, ranks_per_gpu, upper_nnz, tot_local_matrix_nnz, 0,
-        tot_local_matrix_nnz - upper_nnz);
-    auto lower_comm_pattern = compute_gather_to_owner_counts(
-        exec_handler, ranks_per_gpu, upper_nnz, tot_local_matrix_nnz, upper_nnz,
-        nrows + n_interfaces);
+        exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
+        local_matrix_nnz - upper_nnz);
+    auto lower_comm_pattern =
+        compute_gather_to_owner_counts(exec_handler, ranks_per_gpu, upper_nnz,
+                                       local_matrix_nnz, upper_nnz, nrows);
 
     // label nnz = 0;
     //
@@ -176,74 +175,71 @@ void update_fused_impl(
         std::shared_ptr<const LocalMatrixType> mtx;
         // TODO the begin idx need to be exchanged
         label host_interface_ctr{0};
-        label local_begin = local_matrix_nnz;
-        label local_non_local_begin = 0;
-        label non_local_begin = tot_local_matrix_nnz;
-        label non_local_non_local_begin = n_interfaces;
-        label prev_orig_rank = rank;
-        scalar *recv_buffer_ptr;
+
+        label local_elements{0};
+        label non_local_elements{0};
+
+        // compute local_elements and non_local_elements before hand
         for (auto [is_local, orig_rank, size, ctr] : local_interfaces) {
             if (is_local) {
-                mtx =
-                    gko::as<const LocalMatrixType>(dist_A->get_local_matrix());
+                local_elements += size;
             } else {
-                mtx = gko::as<const LocalMatrixType>(
-                    dist_A->get_non_local_matrix());
+                non_local_elements += size;
             }
+        }
+        label local_begin = local_sparsity->num_nnz - local_elements;
+        label non_local_begin = 0;
 
+        scalar *recv_buffer_ptr;
+        for (auto [is_local, orig_rank, size, ctr] : local_interfaces) {
+            mtx =
+                (is_local)
+                    ? gko::as<const LocalMatrixType>(dist_A->get_local_matrix())
+                    : gko::as<const LocalMatrixType>(
+                          dist_A->get_non_local_matrix());
             recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
 
-            label next_begin = size;
-
+            // where to put the data
+            label begin = 0;
+            if (is_local) {
+                begin = local_begin;
+                local_begin += size;
+            } else {
+                begin = non_local_begin;
+                non_local_begin += size;
+            }
             if (orig_rank == rank) {
                 // if data is already on this rank
                 auto data_view = gko::array<scalar>::const_view(
                     exec, size, host_A->get_interface_data(host_interface_ctr));
-
-                label begin = (is_local) ? local_begin : local_non_local_begin;
-
                 auto target_view = gko::array<scalar>::view(
                     mtx->get_executor(), size, recv_buffer_ptr + begin);
                 target_view = data_view;
                 host_interface_ctr++;
-
-                if (is_local) {
-                    local_begin += next_begin;
-                } else {
-                    local_non_local_begin += next_begin;
-                }
-
             } else {
-                // data is not already on rank
-                label begin =
-                    (is_local) ? non_local_begin : non_local_non_local_begin;
-                if (is_local && orig_rank != prev_orig_rank) {
-                    begin += local_matrix_nnz;
-                }
                 comm.recv(device_exec, recv_buffer_ptr + begin, size, orig_rank,
                           0);
-                if (orig_rank == prev_orig_rank) {
-                    if (is_local) {
-                        non_local_begin += next_begin;
-                    } else {
-                        non_local_non_local_begin += next_begin;
-                    }
-                } else {
-                    if (is_local) {
-                        non_local_begin += local_matrix_nnz + next_begin;
-                    } else {
-                        non_local_non_local_begin +=
-                            local_matrix_nnz + next_begin;
-                    }
-                    prev_orig_rank = orig_rank;
-                }
             }
         }
+
         // interface values need to be multiplied by -1
-        using vec = gko::matrix::Dense<scalar>;
+        auto neg_one = gko::initialize<vec>({-1.0}, device_exec);
+        mtx = gko::as<const LocalMatrixType>(dist_A->get_local_matrix());
+        recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
+        label num_interface_elems = tot_local_matrix_nnz - local_matrix_nnz;
+        auto local_interface_dense = vec::create(
+            mtx->get_executor(),
+            gko::dim<2>{static_cast<gko::size_type>(local_elements), 1},
+            gko::array<scalar>::view(
+                mtx->get_executor(), local_elements,
+                recv_buffer_ptr +
+                    (mtx->get_num_stored_elements() - local_elements)),
+            1);
+        local_interface_dense->scale(neg_one);
+
+        // TODO this currently skips scaling of now local interface part
         mtx = gko::as<const LocalMatrixType>(dist_A->get_non_local_matrix());
         recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
-        auto neg_one = gko::initialize<vec>({-1.0}, device_exec);
         auto interface_dense = vec::create(
             mtx->get_executor(),
             gko::dim<2>{
@@ -254,7 +250,7 @@ void update_fused_impl(
             1);
         interface_dense->scale(neg_one);
     } else {
-        // the non owner has send all its interfaces to owner
+        // the non owner has to send all its interfaces to owner
         // thus all values need to be communicated to the owner as well
         label num_interfaces = src_comm_pattern->target_ids.size();
         label tag = 0;
@@ -326,15 +322,15 @@ void update_impl(
         gko::experimental::distributed::Matrix<scalar, label, label>>
         dist_A,
     std::shared_ptr<const SparsityPattern> local_sparsity,
-    [[maybe_unused]] std::shared_ptr<const SparsityPattern> non_local_sparsity,
+    std::shared_ptr<const SparsityPattern> non_local_sparsity,
     std::shared_ptr<const CommunicationPattern> src_comm_pattern,
     std::vector<InterfaceLocality> local_interfaces)
 {
+    using vec = gko::matrix::Dense<scalar>;
     auto exec = exec_handler.get_ref_exec();
     auto device_exec = exec_handler.get_device_exec();
     auto ranks_per_gpu = repartitioner->get_ranks_per_gpu();
-    [[maybe_unused]] bool requires_host_buffer =
-        exec_handler.get_gko_force_host_buffer();
+    bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
 
     label rank{exec_handler.get_rank()};
     label owner_rank = repartitioner->get_owner_rank(exec_handler);
@@ -505,11 +501,10 @@ void update_impl(
                             1},
                 tmp, 1));
 
-        auto non_local_row_vec =
-        gko::share(gko::matrix::Dense<scalar>::create(
+        auto non_local_row_vec = gko::share(gko::matrix::Dense<scalar>::create(
             dist_A->get_executor(),
-            gko::dim<2>{static_cast<dim_type>(non_local_sparsity->num_nnz),
-            1}, gko::array<scalar>(dist_A->get_executor(),
+            gko::dim<2>{static_cast<dim_type>(non_local_sparsity->num_nnz), 1},
+            gko::array<scalar>(dist_A->get_executor(),
                                non_local_sparsity->num_nnz),
             1));
         non_local_sparsity->ldu_mapping.set_executor(dist_A->get_executor());
@@ -525,8 +520,8 @@ void update_impl(
                         dist_A->get_non_local_matrix())
                         ->get_combination()
                         ->get_operators()[i]);
-            auto non_local_elements =
-            non_local_mtx->get_num_stored_elements(); scalar *non_local_ptr =
+            auto non_local_elements = non_local_mtx->get_num_stored_elements();
+            scalar *non_local_ptr =
                 const_cast<scalar *>(non_local_mtx->get_const_values());
             auto value_view =
                 gko::array<scalar>::view(non_local_mtx->get_executor(),
@@ -538,7 +533,7 @@ void update_impl(
             auto non_local_dense_view = gko::array<scalar>::view(
                 dist_A->get_executor(), non_local_elements,
                 non_local_dense_vec->get_values() +
-                non_local_sparsity->spans[i].begin);
+                    non_local_sparsity->spans[i].begin);
             non_local_dense_view = value_view;
         }
 
@@ -546,9 +541,9 @@ void update_impl(
         non_local_dense_vec->row_gather(&non_local_mapping_view,
                                         non_local_row_vec.get());
 
-            auto non_local_dense_view = gko::array<scalar>::view(
-                dist_A->get_executor(), non_local_sparsity->num_nnz,
-                non_local_dense_vec->get_values());
+        auto non_local_dense_view = gko::array<scalar>::view(
+            dist_A->get_executor(), non_local_sparsity->num_nnz,
+            non_local_dense_vec->get_values());
 
         for (size_t i = 0; i < non_local_sparsity->spans.size(); i++) {
             std::shared_ptr<const LocalMatrixType> non_local_mtx =
@@ -557,8 +552,8 @@ void update_impl(
                         dist_A->get_non_local_matrix())
                         ->get_combination()
                         ->get_operators()[i]);
-            auto non_local_elements =
-            non_local_mtx->get_num_stored_elements(); scalar *non_local_ptr =
+            auto non_local_elements = non_local_mtx->get_num_stored_elements();
+            scalar *non_local_ptr =
                 const_cast<scalar *>(non_local_mtx->get_const_values());
             auto value_view =
                 gko::array<scalar>::view(non_local_mtx->get_executor(),
@@ -681,20 +676,6 @@ std::shared_ptr<RepartDistMatrix> create_impl(
         src_comm_pattern, repart_comm_pattern, repartitioner, local_interfaces);
 }
 
-std::shared_ptr<const gko::LinOp> get_local(
-    std::shared_ptr<const gko::LinOp> dist_A)
-{
-    word matrix_format = "Coo";
-    if (matrix_format == "Coo") {
-        return gko::as<RepartDistMatrix>(dist_A)
-            ->get_local<const gko::matrix::Coo<scalar, label>>();
-    }
-    if (matrix_format == "Csr") {
-        return gko::as<RepartDistMatrix>(dist_A)
-            ->get_local<const gko::matrix::Csr<scalar, label>>();
-    }
-}
-
 void write_distributed(const ExecutorHandler &exec_handler, word field_name,
                        const objectRegistry &db,
                        std::shared_ptr<RepartDistMatrix> dist_A,
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index 0df4af496..739156b94 100644
--- a/src/MatrixWrapper/HostMatrix.C
+++ b/src/MatrixWrapper/HostMatrix.C
@@ -302,7 +302,7 @@ HostMatrixWrapper::create_communication_pattern() const
             } else {
                 FatalErrorInFunction
                     << " Currently only unique neighbour ranks are supported "
-                    << " Neighbor proccessor " << neighbProcNo
+                    << " Neighbor processor " << neighbProcNo
                     << " observed twice" << exit(FatalError);
                 // auto &vec = search->second;
                 // vec.insert(vec.end(), face_cells.begin(), face_cells.end());
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index dcf814b92..e103914b3 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -68,8 +68,8 @@ exchange_spans_ranks(const ExecutorHandler &exec_handler, label ranks_per_gpu,
     if (gathered_size.size() > 0) {
         for (int i = 0; i < comm_pattern.recv_counts.size(); i++) {
             label num_interfaces = comm_pattern.recv_counts[i];
-            if ( num_interfaces > 0) {
-                for (int j=0;j<num_interfaces;j++){
+            if (num_interfaces > 0) {
+                for (int j = 0; j < num_interfaces; j++) {
                     origins.push_back(i);
                 }
             }
@@ -118,11 +118,10 @@ label Repartitioner::compute_repart_size(label local_size, label ranks_per_gpu,
     return all_to_all_pattern.recv_offsets.back();
 }
 
+/* reorders the sparsity pattern so that it is in row major order */
 void fuse_sparsity(std::vector<label> &rows, std::vector<label> &cols,
                    std::vector<label> &mapping, std::vector<gko::span> &span)
 {
-    // add offset to mapping
-    // so interface mapping is not continous
     std::vector<label> permutation(rows.size());
     std::iota(permutation.begin(), permutation.end(), 0);
     std::stable_sort(permutation.begin(), permutation.end(),
@@ -142,7 +141,7 @@ void sort_sparsity(std::vector<label> &rows, std::vector<label> &cols,
                    std::vector<label> &mapping, std::vector<gko::span> &span)
 {
     // add offset to mapping
-    // so interface mapping is not continous
+    // so interface mapping is not continuous
     std::vector<label> permutation(rows.size());
     std::iota(permutation.begin(), permutation.end(), 0);
     for (auto [begin, end] : span) {
@@ -225,13 +224,11 @@ Repartitioner::repartition_sparsity(
                 non_local_map, src_non_local_pattern->spans);
         }
 
-	auto copy_local_pattern = std::make_shared<SparsityPattern>(
-                *src_local_pattern.get());
+        auto copy_local_pattern =
+            std::make_shared<SparsityPattern>(*src_local_pattern.get());
 
         LOG_1(verbose_, "done repartition sparsity pattern")
-	return {
-            copy_local_pattern,
-            ret_non_local, std::vector(ret)};
+        return {copy_local_pattern, ret_non_local, std::vector(ret)};
     }
 
     // Step 1. gather all local sparsity pattern to owner rank
@@ -305,16 +302,16 @@ Repartitioner::repartition_sparsity(
     if (is_owner(exec_handler)) {
         LOG_1(verbose_, "done repartition sparsity pattern")
         auto ret = std::make_tuple<std::shared_ptr<SparsityPattern>,
-                               std::shared_ptr<SparsityPattern>,
-                               std::vector<InterfaceLocality>>(
-		std::make_shared<SparsityPattern>(
-            exec, tmp_local_dim, tmp_local_rows, tmp_local_cols,
-            tmp_local_mapping, tmp_local_span),
-           std::make_shared<SparsityPattern>(
-            exec, tmp_non_local_dim, tmp_non_local_rows, tmp_non_local_cols,
-            tmp_non_local_mapping, new_spans)
-            ,std::vector(is_local));
-	return ret;
+                                   std::shared_ptr<SparsityPattern>,
+                                   std::vector<InterfaceLocality>>(
+            std::make_shared<SparsityPattern>(
+                exec, tmp_local_dim, tmp_local_rows, tmp_local_cols,
+                tmp_local_mapping, tmp_local_span),
+            std::make_shared<SparsityPattern>(
+                exec, tmp_non_local_dim, tmp_non_local_rows, tmp_non_local_cols,
+                tmp_non_local_mapping, new_spans),
+            std::vector(is_local));
+        return ret;
     } else {
         LOG_1(verbose_, "done repartition sparsity pattern")
         return std::make_tuple<std::shared_ptr<SparsityPattern>,
@@ -326,8 +323,7 @@ Repartitioner::repartition_sparsity(
 }
 
 
-std::vector<InterfaceLocality>
-Repartitioner::build_non_local_interfaces(
+std::vector<InterfaceLocality> Repartitioner::build_non_local_interfaces(
     const ExecutorHandler &exec_handler,
     std::shared_ptr<
         const gko::experimental::distributed::Partition<label, label>>
@@ -345,7 +341,7 @@ Repartitioner::build_non_local_interfaces(
     std::vector<label> mark_keep;
 
     if (non_local_spans.size() != non_local_rank_origin.size()) {
-                FatalErrorInFunction
+        FatalErrorInFunction
             << "non_local_spans and non_local_rank_origins sizes are different"
             << exit(FatalError);
     }
@@ -358,16 +354,16 @@ Repartitioner::build_non_local_interfaces(
                          return comm_target_ids[i] < comm_target_ids[j];
                      });
 
-    label interface_offset = 0;
+    label interface_offset = local_spans.back().end;
     label ctr = 0;
-    label local_ctr =0;
+    label local_ctr = 0;
     label non_local_ctr = 0;
     for (auto i : iteration_order) {
         auto [begin, end] = non_local_spans[i];
         bool local = reparts_to_local(exec_handler, comm_target_ids[i]);
 
         if (local) {
-            local_ctr++; // local interface starts counting at 1;
+            local_ctr++;  // local interface starts counting at 1;
             gko::size_type rows_start = local_rows.size();
             local_rows.insert(local_rows.end(), non_local_rows.data() + begin,
                               non_local_rows.data() + end);
@@ -392,7 +388,8 @@ Repartitioner::build_non_local_interfaces(
             ctr = non_local_ctr;
             non_local_ctr++;
         }
-        is_local.push_back(InterfaceLocality{local, non_local_rank_origin[i], end - begin, ctr});
+        is_local.push_back(InterfaceLocality{local, non_local_rank_origin[i],
+                                             end - begin, ctr});
     }
 
     // remove data from non_local vectors
@@ -498,10 +495,10 @@ Repartitioner::repartition_comm_pattern(
         for (int i = 1; i < ranks_per_gpu_; i++) {
             // how many gko::arrays to with send_indexes to receive
             // from i-th neighbor
-             label recv_count = comm_pattern.recv_counts[rank + i];
+            label recv_count = comm_pattern.recv_counts[rank + i];
 
             for (int j = 0; j < recv_count; j++) {
-              auto target_size = gathered_target_sizes[j + recv_ctr];
+                auto target_size = gathered_target_sizes[j + recv_ctr];
                 std::vector<label> recv_buffer(target_size);
 
                 comm.recv(exec, recv_buffer.data(), target_size, rank + i,
@@ -519,7 +516,7 @@ Repartitioner::repartition_comm_pattern(
                 // auto target_id = gathered_target_ids[j + owner_recv_counts];
                 send_idxs.emplace_back(recv_buffer);
             }
-            recv_ctr+= recv_count;
+            recv_ctr += recv_count;
         }
     } else {
         label owner = get_owner_rank(exec_handler);
diff --git a/src/StoppingCriterion.C b/src/StoppingCriterion.C
index 75960eb76..dae28c933 100644
--- a/src/StoppingCriterion.C
+++ b/src/StoppingCriterion.C
@@ -97,9 +97,10 @@ bool StoppingCriterion::OpenFOAMDistStoppingCriterion::check_impl(
     auto norm1_host = vec::create(exec->get_master(), gko::dim<2>{1});
     norm1_host->copy_from(norm1.get());
     scalar residual_norm = norm1_host->at(0);
-    if (residual_norm != residual_norm){
-            FatalErrorInFunction << " Problem with residual norm detected: " << residual_norm
-                                 << exit(FatalError);
+    if (residual_norm != residual_norm) {
+        FatalErrorInFunction
+            << " Problem with residual norm detected: " << residual_norm
+            << exit(FatalError);
     }
 
     bool result = false;
diff --git a/src/common.C b/src/common.C
index e4bb77c64..c6ec79792 100644
--- a/src/common.C
+++ b/src/common.C
@@ -176,7 +176,7 @@ std::vector<label> convert_to_vector(const gko::array<label> &in)
 std::vector<scalar> convert_to_vector(const gko::array<scalar> &in)
 {
     return std::vector<scalar>(in.get_const_data(),
-                              in.get_const_data() + in.get_size());
+                               in.get_const_data() + in.get_size());
 }
 
 std::vector<scalar> convert_to_vector(
diff --git a/unitTests/CMakeLists.txt b/unitTests/CMakeLists.txt
index 72916b46a..eb280a0af 100644
--- a/unitTests/CMakeLists.txt
+++ b/unitTests/CMakeLists.txt
@@ -35,6 +35,7 @@ macro(ogl_unit_test TEST FILE)
 endmacro()
 
 ogl_unit_test(CommunicationPattern "CommunicationPattern.C")
-ogl_unit_test(Repartitioner "Repartitioner.C")
+ogl_unit_test(Repartitioner "Repartitioner1D.C")
+ogl_unit_test(Repartitioner2D "Repartitioner2D.C")
 
 add_subdirectory(MatrixWrapper)
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index a2a1dd97c..b87c0e340 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -105,8 +105,6 @@ public:
         }
         for (int i = 0; i < 12; i++) {
             fvMatrix->upper().data()[i] = 1.0;
-        }
-        for (int i = 0; i < 12; i++) {
             fvMatrix->lower().data()[i] = 1.0;
         }
         // set the interface value, we use get_interface_data here
@@ -136,16 +134,8 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new Environment);
 
-class DistributedMatrixFixtureMatrixFormat
-    : public testing::TestWithParam<std::tuple<int, string>> {
-public:
-    ExecutorHandler exec = *((Environment *)global_env)->exec.get();
-    label rank = exec.get_rank();
-    const gko::experimental::mpi::communicator comm =
-        *(exec.get_communicator().get());
-};
-
-class DistributedMatrixFixture : public testing::TestWithParam<int> {
+class DistributedMatrixFixtureFormat
+    : public testing::TestWithParam<std::tuple<int, string, bool>> {
 public:
     ExecutorHandler exec = *((Environment *)global_env)->exec.get();
     label rank = exec.get_rank();
@@ -153,25 +143,280 @@ public:
         *(exec.get_communicator().get());
 
     std::map<label, vec> exp_local_size{
-        {1, {9, 9, 9, 9}}, {2, {18, 0, 18, 0}}, {4, {36, 0, 18, 0}}};
-};
+        {1, {9, 9, 9, 9}}, {2, {18, 0, 18, 0}}, {4, {36, 0, 0, 0}}};
 
+    /*
+     * The mesh has the following structure
+     *         global ids
+     *        [24 25 26|33 34 35]
+     *        [21 22 23|30 31 32]
+     *   2    [18 19 20|27 28 29]  3
+     *        ---------+---------
+     *        [ 6  7  8|15 16 17]
+     *        [ 3  4  5|12 13 14]
+     *   0    [ 0  1  2| 9 10 11]  1
+     *   */
+    std::vector<scalar> exp_local_coeff_1{2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,
+                                          2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2,
+                                          1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
 
-// INSTANTIATE_TEST_SUITE_P(
-//     DistributedMatrixFixtureInstantiationMatrixFormat,
-//     DistributedMatrixFixtureMatrixFormat,
-//     testing::Combine(testing::Values(1,
-//                                      2),  // TODO FIXME case for  4 subdomains
-//                      testing::Values("Coo", "Csr")));
+    std::vector<scalar> exp_local_coeff_2_nf{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1,
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1,
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 3, 1, 2, 3};
 
-INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiation,
-                         DistributedMatrixFixture, testing::Values(1, 2));
+    std::vector<scalar> exp_local_coeff_2_f_1{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,           // 0 - 10
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2,  // 11 - 24
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 3,           // 25 - 35
+        // second element is reported to be wrong
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1,           // 36 - 46
+        2, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,  //
+        3, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
+    std::vector<scalar> exp_local_coeff_2_f_2{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,           //
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2,  //
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 3,           //
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1,           //
+        2, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,  //
+        3, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
+    std::vector<scalar> exp_local_coeff_4{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1,
+        1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1,
+        1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1,
+        1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1,
+        2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 3, 1, 2, 3,
+        1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3};
+    std::vector<scalar> exp_local_coeff_4_f{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,           // 0 - 10
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2,  // 11-24
+        1, 2, 1, 1, 1, 1, 2, 1, 2, 1, 1, 2, 3, 3,  // 25-38
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1,           // 39 - 49
+        2, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,  // 50-64
+        3, 1, 2, 1, 1, 1, 1, 2, 1, 2, 1, 1, 2, 3,  // 65-78
+        1, 2, 1, 1, 2, 1, 2, 1, 1, 3, 1, 2, 1, 1,  // 80-92
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2,  // 93
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 3,           //
+        1, 1, 2, 1, 1, 2, 1, 2, 1, 1, 3, 1, 2, 1,  //
+        2, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,  //
+        3, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
+    std::map<bool, std::map<label, vec_vec_s>> exp_local_coeffs{
+        {true,
+         {{1,
+           {exp_local_coeff_1, exp_local_coeff_1, exp_local_coeff_1,
+            exp_local_coeff_1}},
+          {2, {exp_local_coeff_2_f_1, {}, exp_local_coeff_2_f_2, {}}},
+          {4, {exp_local_coeff_4_f, {}, {}, {}}}}},
+        {false,
+         {{1,
+           {exp_local_coeff_1, exp_local_coeff_1, exp_local_coeff_1,
+            exp_local_coeff_1}},
+          {2, {exp_local_coeff_2_nf, {}, exp_local_coeff_2_nf, {}}},
+          {4, {exp_local_coeff_4, {}, {}, {}}}}}};
+
+    std::vector<scalar> non_local_coeffs{1, 2, 3, 1, 2, 3};
+    std::map<bool, std::map<label, vec_vec_s>> exp_non_local_coeffs{
+        {true,
+         {{1,
+           {{1, 2, 1, 2, 3, 3},
+            {1, 2, 3, 1, 2, 3},
+            {1, 2, 3, 1, 2, 3},
+            {1, 1, 2, 3, 2, 3}}},
+          {2, {non_local_coeffs, {}, non_local_coeffs, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {false,
+         {{1,
+           {non_local_coeffs, non_local_coeffs, non_local_coeffs,
+            non_local_coeffs}},
+          {2, {{1, 2, 3, 1, 2, 3}, {}, {1, 2, 3, 1, 2, 3}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_rows{
+        {true,
+         {{1,
+           {{2, 5, 6, 7, 8, 8},
+            {0, 3, 6, 6, 7, 8},
+            {0, 1, 2, 2, 5, 8},
+            {0, 0, 1, 2, 3, 6}}},
+          {2, {{6, 7, 8, 15, 16, 17}, {}, {0, 1, 2, 9, 10, 11}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {false,
+         {{1,
+           {{2, 5, 8, 6, 7, 8},
+            {0, 3, 6, 6, 7, 8},
+            {0, 1, 2, 2, 5, 8},
+            {0, 1, 2, 0, 3, 6}}},
+          {2, {{6, 7, 8, 15, 16, 17}, {}, {0, 1, 2, 9, 10, 11}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_cols{
+        {true,
+         {{1,
+           {{0, 1, 3, 4, 2, 5},
+            {0, 1, 2, 3, 4, 5},
+            {0, 1, 2, 3, 4, 5},
+            {0, 3, 1, 2, 4, 5}}},
+          {2, {{0, 1, 2, 3, 4, 5}, {}, {0, 1, 2, 3, 4, 5}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {false,
+         {{1,
+           {{0, 1, 2, 3, 4, 5},
+            {0, 1, 2, 3, 4, 5},
+            {0, 1, 2, 3, 4, 5},
+            {0, 1, 2, 3, 4, 5}}},
+          {2, {{0, 1, 2, 3, 4, 5}, {}, {0, 1, 2, 3, 4, 5}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
 
-TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
+    /*
+     * The mesh has the following structure
+     *         local ids
+     *        [24 25 26|33 34 35]
+     *        [21 22 23|30 31 32]
+     *        [18 19 20|27 28 29]
+     *        ---------+---------
+     *        [ 6  7  8|15 16 17]
+     *        [ 3  4  5|12 13 14]
+     *   0    [ 0  1  2| 9 10 11]  1
+     *   */
+    vec local_row_1 = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4,
+                       4, 4, 5, 5, 5, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8};
+    vec local_row_2 = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,  4,
+        4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  9,  9,  9,
+        10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12, 13, 13, 13, 13, 13, 14, 14,
+        14, 14, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17, 9,  12, 15, 2,  5,  8};
+    vec local_row_2_f = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,
+        4,  4,  5,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  8,
+        9,  9,  9,  9,  10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12, 12, 13, 13,
+        13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17};
+    vec local_row_4 = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,  4,
+        4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  9,  9,  9,
+        10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12, 13, 13, 13, 13, 13, 14, 14,
+        14, 14, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
+        19, 20, 20, 20, 21, 21, 21, 21, 22, 22, 22, 22, 22, 23, 23, 23, 23, 24,
+        24, 24, 25, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 28, 29, 29,
+        29, 30, 30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32, 32, 33, 33, 33, 34,
+        34, 34, 34, 35, 35, 35, 9,  12, 15, 18, 19, 20, 2,  5,  8,  27, 28, 29,
+        6,  7,  8,  27, 30, 33, 15, 16, 17, 20, 23, 26};
+    vec local_row_4_f = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,
+        4,  4,  5,  5,  5,  5,  5,  6,  6,  6,  6,  7,  7,  7,  7,  7,  8,  8,
+        8,  8,  8,  9,  9,  9,  9,  10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12,
+        12, 13, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 15, 15, 16, 16, 16,
+        16, 16, 17, 17, 17, 17, 18, 18, 18, 18, 19, 19, 19, 19, 19, 20, 20, 20,
+        20, 20, 21, 21, 21, 21, 22, 22, 22, 22, 22, 23, 23, 23, 23, 23, 24, 24,
+        24, 25, 25, 25, 25, 26, 26, 26, 26, 27, 27, 27, 27, 27, 28, 28, 28, 28,
+        28, 29, 29, 29, 29, 30, 30, 30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32,
+        32, 33, 33, 33, 33, 34, 34, 34, 34, 35, 35, 35};
+    std::map<bool, std::map<label, vec_vec>> exp_local_rows{
+        {true,
+         {{1, {local_row_1, local_row_1, local_row_1, local_row_1}},
+          {2, {local_row_2_f, {}, local_row_2_f, {}}},
+          {4, {local_row_4_f, {}, {}, {}}}}},
+        {false,
+         {{1, {local_row_1, local_row_1, local_row_1, local_row_1}},
+          {2, {local_row_2, {}, local_row_2, {}}},
+          {4, {local_row_4, {}, {}, {}}}}}};
+
+    vec local_cols_1 = {0, 1, 3, 0, 1, 2, 4, 1, 2, 5, 0, 3, 4, 6, 1, 3, 4,
+                        5, 7, 2, 4, 5, 8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8};
+    vec local_cols_2 = {
+        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
+        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
+        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
+        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 2,  5,  8,  9,  12, 15};
+    vec local_cols_2_f = {
+        0, 1,  3,  0,  1,  2,  4,  1,  2,  5,  9,               //
+        0, 3,  4,  6,  1,  3,  4,  5,  7,  2,  4,  5,  8,  12,  //
+        3, 6,  7,  4,  6,  7,  8,  5,  7,  8,  15,              //
+        2, 9,  10, 12, 9,  10, 11, 13, 10, 11, 14,              //
+        5, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13, 14, 17,  //
+        8, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17};
+    vec local_cols_4 = {
+        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
+        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
+        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
+        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 18, 19, 21, 18, 19, 20,
+        22, 19, 20, 23, 18, 21, 22, 24, 19, 21, 22, 23, 25, 20, 22, 23, 26, 21,
+        24, 25, 22, 24, 25, 26, 23, 25, 26, 27, 28, 30, 27, 28, 29, 31, 28, 29,
+        32, 27, 30, 31, 33, 28, 30, 31, 32, 34, 29, 31, 32, 35, 30, 33, 34, 31,
+        33, 34, 35, 32, 34, 35, 2,  5,  8,  6,  7,  8,  9,  12, 15, 15, 16, 17,
+        18, 19, 20, 20, 23, 26, 27, 28, 29, 27, 30, 33};
+    vec local_cols_4_f = {
+        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  9,               // 0 - 10
+        0,  3,  4,  6,  1,  3,  4,  5,  7,  2,  4,  5,  8,  12,  // 11 - 24
+        3,  6,  7,  18, 4,  6,  7,  8,  19, 5,  7,  8,  15, 20,  // 12 - 38
+        2,  9,  10, 12, 9,  10, 11, 13, 10, 11, 14,              //
+        5,  9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13, 14, 17,  //
+        8,  12, 15, 16, 27, 13, 15, 16, 17, 28, 14, 16, 17, 29,  //
+        6,  18, 19, 21, 7,  18, 19, 20, 22, 8,  19, 20, 23, 27,  //
+        18, 21, 22, 24, 19, 21, 22, 23, 25, 20, 22, 23, 26, 30,  //
+        21, 24, 25, 22, 24, 25, 26, 23, 25, 26, 33,              //
+        15, 20, 27, 28, 30, 16, 27, 28, 29, 31, 17, 28, 29, 32,  //
+        23, 27, 30, 31, 33, 28, 30, 31, 32, 34, 29, 31, 32, 35,  //
+        26, 30, 33, 34, 31, 33, 34, 35, 32, 34, 35};
+    std::map<bool, std::map<label, vec_vec>> exp_local_cols{
+        {true,
+         {{1, {local_cols_1, local_cols_1, local_cols_1, local_cols_1}},
+          {2, {local_cols_2_f, {}, local_cols_2_f, {}}},
+          {4, {local_cols_4_f, {}, {}, {}}}}},
+        {false,
+         {{1, {local_cols_1, local_cols_1, local_cols_1, local_cols_1}},
+          {2, {local_cols_2, {}, local_cols_2, {}}},
+          {4, {local_cols_4, {}, {}, {}}}}}};
+
+    vec_vec_s x_1{{4, 5, 5, 5, 6, 7, 5, 7, 10},
+                  {5, 5, 4, 7, 6, 5, 8, 7, 7},
+                  {5, 7, 8, 5, 6, 7, 4, 5, 7},
+                  {6, 7, 7, 7, 6, 5, 7, 5, 4}};
+    std::vector<scalar> exp_x_2_1 = {4, 5, 5, 5, 6, 7, 5, 7, 10,
+                                     5, 5, 4, 7, 6, 5, 8, 7, 7};
+    std::vector<scalar> exp_x_2_2 = {5, 7, 8, 5, 6, 7, 4, 5, 7,
+                                     6, 7, 7, 7, 6, 5, 7, 5, 4};
+    std::vector<scalar> exp_x_4 = {4, 5, 5, 5, 6, 7, 5, 7, 10, 5, 5, 4,
+                                   7, 6, 5, 8, 7, 7, 5, 7, 8,  5, 6, 7,
+                                   4, 5, 7, 6, 7, 7, 7, 6, 5,  7, 5, 4};
+    std::map<bool, std::map<label, vec_vec_s>> exp_x{
+        {true,
+         {{1, x_1},
+          {2, {exp_x_2_1, {}, exp_x_2_2, {}}},
+          {4, {exp_x_4, {}, {}, {}}}}},
+        {false,
+         {{1, x_1},
+          {2, {exp_x_2_1, {}, exp_x_2_2, {}}},
+          {4, {exp_x_4, {}, {}, {}}}}}};
+};
+
+
+INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiationMatrixFormat,
+                         DistributedMatrixFixtureFormat,
+                         testing::Combine(testing::Values(1, 2, 4),
+                                          testing::Values("Coo"),
+                                          testing::Values(false, true)),
+                         [](const auto &info) {
+                             // Can use info.param here to generate the test
+                             // suffix
+                             std::vector<std::string> names;
+                             names.emplace_back("ranks");
+                             names.emplace_back("format");
+                             names.emplace_back("fuse");
+                             std::string name = "ranks_";
+                             name += std::to_string(std::get<0>(info.param));
+                             name += "_format_";
+                             name += std::get<1>(info.param);
+                             name += "_fused_";
+                             name += std::to_string(std::get<2>(info.param));
+                             return name;
+                         });
+
+
+TEST_P(DistributedMatrixFixtureFormat, canCreateDistributedMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    auto ranks_per_gpu = GetParam();
-    bool fused = false;
+    auto [ranks_per_gpu, matrix_format, fused] = GetParam();
 
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
@@ -183,7 +428,7 @@ TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
     gko::dim<2> local_vec_dim{repartitioner->get_repart_dim()[0], 1};
 
     auto distributed =
-        create_distributed(exec, repartitioner, hostMatrix, "Coo");
+        create_distributed(exec, repartitioner, hostMatrix, matrix_format);
 
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
               exp_local_size[ranks_per_gpu][rank]);
@@ -195,11 +440,10 @@ TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
               exp_local_size[ranks_per_gpu][rank]);
 }
 
-TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
+TEST_P(DistributedMatrixFixtureFormat, hasCorrectLocalMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    auto ranks_per_gpu = GetParam();
-    bool fused = false;
+    auto [ranks_per_gpu, matrix_format, fused] = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner = std::make_shared<Repartitioner>(
@@ -209,115 +453,52 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
                                1};
     gko::dim<2> local_vec_dim{repartitioner->get_repart_dim()[0], 1};
 
-
     auto distributed =
-        create_distributed(exec, repartitioner, hostMatrix, "Coo");
-
-    auto local = detail::convert_combination_to_coo(
-        exec.get_ref_exec(), distributed->get_local_matrix());
-
-    auto res_local_values = convert_to_vector(get_val(local));
-    std::map<label, vec_vec_s> exp_local_values;
-    std::vector<scalar> exp_val_1{2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,
-                                  2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2,
-                                  1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
-    exp_local_values.emplace(
-        1, vec_vec_s{exp_val_1, exp_val_1, exp_val_1, exp_val_1});
-    std::vector<scalar> exp_val_2{
-        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1,
-        2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1,
-        1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 3, 1, 2, 3};
-    exp_local_values.emplace(2, vec_vec_s{exp_val_2, {}, exp_val_2, {}});
-    std::vector<scalar> exp_val_4{
-        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,
-        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1,
-        1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1,
-        1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1,
-        1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1,
-        2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 2, 3, 1, 2, 3,
-        1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3};
-    exp_local_values.emplace(4, vec_vec_s{exp_val_4, {}, {}, {}});
-
-    auto res_local_cols = convert_to_vector(get_col(local));
-    std::map<label, vec_vec> exp_local_cols;
-    vec exp_cols_1 = {0, 1, 3, 0, 1, 2, 4, 1, 2, 5, 0, 3, 4, 6, 1, 3, 4,
-                      5, 7, 2, 4, 5, 8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8};
-    exp_local_cols.emplace(
-        1, vec_vec{exp_cols_1, exp_cols_1, exp_cols_1, exp_cols_1});
-    vec exp_cols_2 = {
-        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
-        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
-        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
-        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 2,  5, 8, 9,  12,  15};
-    exp_local_cols.emplace(2, vec_vec{exp_cols_2, {}, exp_cols_2, {}});
-    vec exp_col_4 = {
-        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
-        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
-        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
-        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 18, 19, 21, 18, 19, 20,
-        22, 19, 20, 23, 18, 21, 22, 24, 19, 21, 22, 23, 25, 20, 22, 23, 26, 21,
-        24, 25, 22, 24, 25, 26, 23, 25, 26, 27, 28, 30, 27, 28, 29, 31, 28, 29,
-        32, 27, 30, 31, 33, 28, 30, 31, 32, 34, 29, 31, 32, 35, 30, 33, 34, 31,
-        33, 34, 35, 32, 34, 35, 2,  5, 8, 6, 7, 8, 9,  12,  15,  15, 16, 17,
-        18,  19,  20,  20, 23, 26, 27, 28, 29, 27, 30, 33};
-    exp_local_cols.emplace(4, vec_vec{exp_col_4, {}, {}, {}});
+        create_distributed(exec, repartitioner, hostMatrix, matrix_format);
 
-    auto res_local_rows = convert_to_vector(get_row(local));
-    std::map<label, vec_vec> exp_local_rows;
-    vec exp_row_1 = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4,
-                     4, 4, 5, 5, 5, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8};
-    exp_local_rows.emplace(1,
-                           vec_vec{exp_row_1, exp_row_1, exp_row_1, exp_row_1});
-    /*
-     * The mesh has the following structure
-     *         local ids
-     *        -------+---------
-     *        [ 6 7 8|15 16 17]
-     *        [ 3 4 5|12 13 14]
-     *   0    [ 0 1 2| 9 10 11]  1
-     *   */
-    vec exp_row_2 = {0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,
-                     4,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,
-                     8,  8,  8,  9,  9,  9,  10, 10, 10, 10, 11, 11, 11, 12, 12,
-                     12, 12, 13, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 16,
-                     16, 16, 16, 17, 17, 17, 9,  12,  15,  2,  5, 8};
-    exp_local_rows.emplace(2, vec_vec{exp_row_2, {}, exp_row_2, {}});
-    /*
-     * The mesh has the following structure
-     *         local ids
-     *        [24 25 26|33 34 35
-     *        [21 22 23|30 31 32
-     *        [18 19 20|27 28 29
-     *        ---------+---------
-     *        [ 6  7  8|15 16 17]
-     *        [ 3  4  5|12 13 14]
-     *   0    [ 0  1  2| 9 10 11]  1
-     *   */
-    vec exp_row_4 = {
-        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,  4,
-        4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  9,  9,  9,
-        10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12, 13, 13, 13, 13, 13, 14, 14,
-        14, 14, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
-        19, 20, 20, 20, 21, 21, 21, 21, 22, 22, 22, 22, 22, 23, 23, 23, 23, 24,
-        24, 24, 25, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 28, 29, 29,
-        29, 30, 30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32, 32, 33, 33, 33, 34,
-        34, 34, 34, 35, 35, 35, 9,  12,  15,  18,  19,  20,  2,  5, 8, 27, 28, 29,
-        6, 7, 8, 27, 30, 33, 15, 16, 17, 20, 23, 26};
-    exp_local_rows.emplace(4, vec_vec{exp_row_4, {}, {}, {}});
+    auto local =
+        (fused) ? gko::as<gko::matrix::Coo<scalar, label>>(
+                      distributed->get_local_matrix())
+                : detail::convert_combination_to_coo(
+                      exec.get_ref_exec(), distributed->get_local_matrix());
 
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[1],
               exp_local_size[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_values, exp_local_values[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_cols, exp_local_cols[ranks_per_gpu][rank]);
+
+    auto res_local_coeffs = convert_to_vector(get_val(local));
+    auto res_local_cols = convert_to_vector(get_col(local));
+    auto res_local_rows = convert_to_vector(get_row(local));
+
+    EXPECT_EQ(res_local_rows.size(),
+              exp_local_rows[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_rows.size(); i++) {
+        ASSERT_EQ(res_local_rows[i],
+                  exp_local_rows[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i << " on rank " << rank;
+    }
+
+    EXPECT_EQ(res_local_cols.size(),
+              exp_local_cols[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_rows.size(); i++) {
+        ASSERT_EQ(res_local_cols[i],
+                  exp_local_cols[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i << " on rank " << rank;
+    }
+
+    EXPECT_EQ(res_local_coeffs.size(),
+              exp_local_coeffs[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_rows.size(); i++) {
+        ASSERT_EQ(res_local_coeffs[i],
+                  exp_local_coeffs[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i << " on rank " << rank;
+    }
 }
 
 
-TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
+TEST_P(DistributedMatrixFixtureFormat, hasCorrectNonLocalMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    auto ranks_per_gpu = GetParam();
-    bool fused = false;
+    auto [ranks_per_gpu, matrix_format, fused] = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner = std::make_shared<Repartitioner>(
@@ -329,98 +510,60 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
     exp_non_local_size.emplace(4, vec{0, 0, 0, 0});
 
     auto distributed =
-        create_distributed(exec, repartitioner, hostMatrix, "Coo");
-
-    auto non_local = detail::convert_combination_to_coo(
-        exec.get_ref_exec(), distributed->get_non_local_matrix());
+        create_distributed(exec, repartitioner, hostMatrix, matrix_format);
 
-    auto res_non_local_values = convert_to_vector(get_val(non_local));
-
-    std::map<label, vec_vec_s> exp_non_local_values;
-    std::vector<scalar> exp_non_local {1.0, 2.0, 3.0, 1.0, 2.0, 3.0};
-    exp_non_local_values.emplace(1, vec_vec_s(4, exp_non_local));
-    exp_non_local_values.emplace(
-        2, vec_vec_s{exp_non_local, {}, exp_non_local, {}});
-    exp_non_local_values.emplace(4, vec_vec_s{{}, {}, {}, {}});
+    auto non_local =
+        (fused) ? gko::as<gko::matrix::Coo<scalar, label>>(
+                      distributed->get_non_local_matrix())
+                : detail::convert_combination_to_coo(
+                      exec.get_ref_exec(), distributed->get_non_local_matrix());
 
+    auto res_non_local_coeffs = convert_to_vector(get_val(non_local));
     auto res_non_local_cols = convert_to_vector(get_col(non_local));
-    std::map<label, vec_vec> exp_non_local_cols;
-    exp_non_local_cols.emplace(1, vec_vec{{0, 1, 2, 3, 4, 5},
-                                          {0, 1, 2, 3, 4, 5},
-                                          {0, 1, 2, 3, 4, 5},
-                                          {0, 1, 2, 3, 4, 5}});
-    exp_non_local_cols.emplace(
-        2, vec_vec{{0, 1, 2, 3, 4, 5}, {}, {0, 1, 2, 3, 4, 5}, {}});
-    exp_non_local_cols.emplace(4, vec_vec{{}, {}, {}, {}});
-
     auto res_non_local_rows = convert_to_vector(get_row(non_local));
-    std::map<label, vec_vec> exp_non_local_rows;
-    exp_non_local_rows.emplace(1, vec_vec{{2, 5, 8, 6, 7, 8},
-                                          {0, 3, 6, 6, 7, 8},
-                                          {0, 1, 2, 2, 5, 8},
-                                          {0, 1, 2, 0, 3, 6}});
-    exp_non_local_rows.emplace(
-        2, vec_vec{{6, 7, 8, 15, 16, 17}, {}, {0, 1, 2, 9, 10, 11}, {}});
-    exp_non_local_rows.emplace(4, vec_vec{{}, {}, {}, {}});
 
     ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[1],
               exp_non_local_size[ranks_per_gpu][rank]);
 
-    ASSERT_EQ(res_non_local_values, exp_non_local_values[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_non_local_rows, exp_non_local_rows[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_non_local_cols, exp_non_local_cols[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_coeffs,
+              exp_non_local_coeffs[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_rows,
+              exp_non_local_rows[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_cols,
+              exp_non_local_cols[fused][ranks_per_gpu][rank]);
 }
 
-// TEST_P(DistributedMatrixFixtureMatrixFormat, distributedMatrixCanApplyCorrectly)
-// {
-//     auto [ranks_per_gpu, format] = GetParam();
-//     bool fused = false;
-//     auto mesh = ((Environment *)global_env)->mesh;
-//     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
-//     auto repartitioner = std::make_shared<Repartitioner>(
-//         hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec, fused);
-
-//     auto distributed =
-//         create_distributed(exec, repartitioner, hostMatrix, format);
-
-//     gko::dim<2> global_vec_dim{repartitioner->get_orig_partition()->get_size(),
-//                                1};
-//     gko::dim<2> local_vec_dim{repartitioner->get_repart_dim()[0], 1};
-
-//     auto b = gko::share(gko::experimental::distributed::Vector<scalar>::create(
-//         exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
-//     b->fill(1);
-
-//     auto x = gko::share(gko::experimental::distributed::Vector<scalar>::create(
-//         exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
-//     x->fill(0);
-
-//     std::map<int, vec_vec_s> exp_x_local;
-
-//     exp_x_local.emplace(1, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6},
-//                                      {5, 5, 4, 6, 6, 5, 6, 6, 5},
-//                                      {5, 6, 6, 5, 6, 6, 4, 5, 5},
-//                                      {6, 6, 5, 6, 6, 5, 5, 5, 4}});
-
-//     std::vector<scalar> exp_x_local_2_1 = {4, 5, 5, 5, 6, 6, 5, 6, 6,
-//                                            5, 5, 4, 6, 6, 5, 6, 6, 5};
-//     std::vector<scalar> exp_x_local_2_2 = {5, 6, 6, 5, 6, 6, 4, 5, 5,
-//                                            6, 6, 5, 6, 6, 5, 5, 5, 4};
-//     exp_x_local.emplace(2, vec_vec_s{exp_x_local_2_1, {}, exp_x_local_2_2, {}});
-
-//     std::vector<scalar> exp_x_local_4 = {4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4,
-//                                          6, 6, 5, 6, 6, 5, 5, 6, 6, 5, 6, 6,
-//                                          4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4};
-//     exp_x_local.emplace(4, vec_vec_s{exp_x_local_4, {}, {}, {}});
-
-//     // Act
-//     distributed->apply(b, x);
-//     auto x_local = std::vector<scalar>(
-//         x->get_local_vector()->get_const_values(),
-//         x->get_local_vector()->get_const_values() + local_vec_dim[0]);
-
-//     ASSERT_EQ(x_local, exp_x_local[ranks_per_gpu][rank]);
-// }
+TEST_P(DistributedMatrixFixtureFormat, canApplyCorrectly)
+{
+    auto [ranks_per_gpu, format, fused] = GetParam();
+    auto mesh = ((Environment *)global_env)->mesh;
+    auto hostMatrix = ((Environment *)global_env)->hostMatrix;
+    auto repartitioner = std::make_shared<Repartitioner>(
+        hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec, fused);
+
+    auto distributed =
+        create_distributed(exec, repartitioner, hostMatrix, format);
+
+    gko::dim<2> global_vec_dim{repartitioner->get_orig_partition()->get_size(),
+                               1};
+    gko::dim<2> local_vec_dim{repartitioner->get_repart_dim()[0], 1};
+
+    auto b = gko::share(gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
+    b->fill(1);
+
+    auto x = gko::share(gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
+    x->fill(0);
+
+    // Act
+    distributed->apply(b, x);
+    auto res_x = std::vector<scalar>(
+        x->get_local_vector()->get_const_values(),
+        x->get_local_vector()->get_const_values() + local_vec_dim[0]);
+
+    ASSERT_EQ(res_x, exp_x[fused][ranks_per_gpu][rank]);
+}
 
 int main(int argc, char *argv[])
 {
diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index cebf233c3..d43c8b483 100644
--- a/unitTests/MatrixWrapper/HostMatrix.C
+++ b/unitTests/MatrixWrapper/HostMatrix.C
@@ -204,7 +204,7 @@ TEST(HostMatrix, canGenerateLocalSparsityPattern)
     //     8,  10, 19, 11,     // cell 7
     //     9,  11, 20          // cell 8
     // });
-    // asymetric case
+    // asymmetric case
     std::vector<label> mapping_expected({
         24, 0,  1,          // cell 0
         12, 25, 2,  3,      // cell 1
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
deleted file mode 100644
index 7ec97fbfd..000000000
--- a/unitTests/Repartitioner.C
+++ /dev/null
@@ -1,863 +0,0 @@
-// SPDX-FileCopyrightText: 2024 OGL authors
-//
-// SPDX-License-Identifier: GPL-3.0-or-later
-
-#include "OGL/Repartitioner.H"
-#include "OGL/common.H"
-
-
-#include "gtest/gtest.h"
-
-#include "mpi.h"
-
-#include <cstdlib>
-
-//---------------------------------------------
-// some_header.h
-extern int my_argc;
-extern char **my_argv;
-// eof
-//---------------------------------------------
-
-//---------------------------------------------
-// main.cpp
-int my_argc;
-char **my_argv;
-
-using vec = std::vector<label>;
-using vec_vec = std::vector<std::vector<label>>;
-using vec_vec_vec = std::vector<vec_vec>;
-
-class RepartitionerEnvironment : public testing::Environment {
-public:
-    void SetUp()
-    {
-        args_ = std::make_shared<Foam::argList>(my_argc, my_argv);
-        time = std::make_shared<Foam::Time>("controlDict", *args_.get());
-        db = std::make_shared<Foam::objectRegistry>(*time.get());
-        Foam::dictionary dict;
-        dict.add("executor", "reference");
-        exec = std::make_shared<ExecutorHandler>(time->thisDb(), dict, "dummy",
-                                                 true);
-
-        if (exec->get_communicator()->size() < 2) {
-            std::cout << "At least 2 CPU processes should be used!"
-                      << std::endl;
-            std::abort();
-        }
-        // delete listener on ranks != 0
-        // to clean up output
-        ::testing::TestEventListeners &listeners =
-            ::testing::UnitTest::GetInstance()->listeners();
-        // if (exec->get_communicator()->rank() != 0) {
-        //     delete listeners.Release(listeners.default_result_printer());
-        // }
-    }
-
-    std::shared_ptr<Foam::argList> args_;
-    std::shared_ptr<const Foam::Time> time;
-    std::shared_ptr<Foam::objectRegistry> db;
-    std::shared_ptr<ExecutorHandler> exec;
-};
-
-const testing::Environment *global_env =
-    AddGlobalTestEnvironment(new RepartitionerEnvironment);
-
-class RepartitionerFixture : public testing::Test {
-public:
-    ExecutorHandler exec =
-        *((RepartitionerEnvironment *)global_env)->exec.get();
-    label rank = exec.get_rank();
-    const gko::experimental::mpi::communicator comm =
-        *(exec.get_communicator().get());
-};
-
-/* @brief Test fixture class for 1D mesh
- *
- * The mesh has the following structure
- * ranks:    0     1     2     3
- * cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
- * cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
- */
-class RepartitionerFixture1D
-    : public RepartitionerFixture,
-      public testing::WithParamInterface<std::tuple<bool, int>> {
-public:
-    label local_size = 2;  // matrix rows, same for all ranks
-
-    // local data
-    std::vector<label> rows{0, 0, 1, 1};
-    std::vector<label> cols{0, 1, 0, 1};
-    std::vector<label> mapping{0, 1, 2, 3};
-    std::vector<gko::span> spans{gko::span{0, 5}};
-
-    // Setup data
-    // non local data
-    vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
-    // non local columns in global idxs
-    vec_vec non_local_cols{{2}, {1, 4}, {3, 6}, {5}};
-    vec_vec non_local_mapping{{0}, {0, 1}, {0, 1}, {0}};
-    // communication partners (ranks)
-    vec_vec ids{{1}, {0, 2}, {1, 3}, {2}};
-
-    std::vector<std::vector<gko::span>> non_local_spans{
-        {gko::span{0, 1}},                   // rank 0
-        {gko::span{0, 1}, gko::span{1, 2}},  // rank 1
-        {gko::span{0, 1}, gko::span{1, 2}},  // rank 2
-        {gko::span{0, 1}},                   // rank 3
-    };
-
-    vec_vec non_local_ranks{
-        {1},     // rank 0
-        {0, 2},  // rank 1
-        {1, 3},  // rank 2
-        {2},     // rank 3
-    };
-
-    // Expected results
-    // number of non-zeros of each sparsity pattern
-    std::map<label, std::vector<label>> exp_local_nnz{
-        {1, {4, 4, 4, 4}},
-        {2, {10, 0, 10, 0}},
-        {4, {22, 0, 0, 0}},
-    };
-
-    std::map<label, std::vector<label>> exp_non_local_nnz{
-        {1, {1, 2, 2, 1}},
-        {2, {1, 0, 1, 0}},
-        {4, {0, 0, 0, 0}},
-    };
-
-    std::vector<std::vector<label>> comm_target_ids{{1}, {0, 2}, {1, 3}, {2}};
-    // expected local row indices in local indices
-    // first map fused true/false
-    // local row indices
-    // repartitioned 2 [ 0 1 , 2 3 | 0 1 , 2 3 ] | new  boundary , old boundary
-    // repartitioned 4 [ 0 1 , 2 3 , 4 5 , 6 7 ] | new  boundary , old boundary
-    // the last two elements are now local interfaces, they are in in order
-    // of target_ids
-    vec nf_rows_2 = {0, 0, 1, 1, 2, 2, 3, 3, 2, 1};
-    vec nf_rows_4 = {0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5,
-                     5, 6, 6, 7, 7, 2, 1, 4, 3, 6, 5};
-    // fused
-    vec f_rows_2 = {0, 0, 1, 1, 1, 2, 2, 2, 3, 3};
-    vec f_rows_4 = {0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3,
-                    4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7};
-    std::map<bool, std::map<label, vec_vec>> exp_local_rows{
-        {false,
-         {{1, {rows, rows, rows, rows}},
-          {2, {nf_rows_2, {}, nf_rows_2, {}}},
-          {4, {nf_rows_4, {}, {}, {}}}}},
-        {true,
-         {{1, {rows, rows, rows, rows}},
-          {2, {f_rows_2, {}, f_rows_2, {}}},
-          {4, {f_rows_4, {}, {}, {}}}}}};
-
-    vec nf_cols_2 = {0, 1, 0, 1, 2, 3, 2, 3, 1, 2};
-    vec nf_cols_4 = {0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4,
-                     5, 6, 7, 6, 7, 1, 2, 3, 4, 5, 6};
-    // fused
-    vec f_cols_2 = {0, 1, 0, 1, 2, 1, 2, 3, 2, 3};
-    vec f_cols_4 = {0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4,
-                    3, 4, 5, 4, 5, 6, 5, 6, 7, 6, 7};
-    // first map fused true/false
-    // expected local col indices in local indices
-    std::map<bool, std::map<label, vec_vec>> exp_local_cols{
-        {false,
-         {{1, {cols, cols, cols, cols}},
-          {2, {nf_cols_2, {}, nf_cols_2, {}}},
-          {4, {nf_cols_4, {}, {}, {}}}}},
-        {true,
-         {{1, {cols, cols, cols, cols}},
-          {2, {f_cols_2, {}, f_cols_2, {}}},
-          {4, {f_cols_4, {}, {}, {}}}}}};
-
-    vec nf_map_2{0, 1, 2, 3, 4, 5, 6, 7,  // here the interface values start
-                 0, 1};                   // <- they are currently unused
-    vec nf_map_4{0, 1, 2,  3,  4,  5,  6,  7,
-                 8, 9, 10, 11, 12, 13, 14, 15,  // here the interface start
-                 0, 1, 2,  3,  4,  5};          // <- unused values because we
-                                                // dont permute atm
-    vec f_map_2{0, 1, 2, 3, 8, 9, 4, 5, 6, 7};
-    vec f_map_4{0,  1, 2, 3,  16, 17, 4,  5,  6,  7,  18,
-                19, 8, 9, 10, 11, 20, 21, 12, 13, 14, 15};
-    std::map<bool, std::map<label, vec_vec>> exp_local_mapping{
-        {false,
-         {{1, {mapping, mapping, mapping, mapping}},
-          {2, {nf_map_2, {}, nf_map_2, {}}},
-          {4, {nf_map_4, {}, {}, {}}}}},
-        {true,
-         {{1, {mapping, mapping, mapping, mapping}},
-          {2, {f_map_2, {}, f_map_2, {}}},
-          {4, {f_map_4, {}, {}, {}}}}}};
-
-    // [start,end) indices of interface/submatrices
-    std::map<bool, std::map<label, vec_vec>> exp_local_spans_begin{
-        {false,
-         {{1, {{0}, {0}, {0}, {0}}},
-          {2, {{0, 8, 9}, {}, {0, 8, 9}, {}}},
-          {4, {{0, 16, 17, 18, 19, 20, 21}, {}, {}, {}}}}},
-        {true,
-         {{1, {{0}, {0}, {0}, {0}}},
-          {2, {{0}, {}, {0}, {}}},
-          {4, {{0}, {}, {}, {}}}}}};
-
-    // [start,end) indices of interface/submatrices
-    std::map<bool, std::map<label, vec_vec>> exp_local_spans_end{
-        {false,
-         {{1, {{5}, {5}, {5}, {5}}},
-          {2, {{8, 9, 10}, {}, {8, 9, 10}, {}}},
-          {4, {{16, 17, 18, 19, 20, 21, 22}, {}, {}, {}}}}},
-        {true,
-         {{1, {{5}, {5}, {5}, {5}}},
-          {2, {{10}, {}, {10}, {}}},
-          {4, {{22}, {}, {}, {}}}}}};
-
-    // number of rows of local matrix
-    std::map<label, vec> exp_local_dim{
-        {1, {local_size, local_size, local_size, local_size}},
-        {2, {2 * local_size, 0, 2 * local_size, 0}},
-        {4, {4 * local_size, 0, 0, 0}},
-    };
-
-    // non local data
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_rows{
-        {false,
-         {{1, non_local_rows}, {2, {{3}, {}, {0}, {}}}, {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, non_local_rows},
-          {2, {{3}, {}, {0}, {}}},
-          {4, {{}, {}, {}, {}}}}}};
-
-    // non local cols are in global indices
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_cols{
-        {false,
-         {{1, non_local_cols}, {2, {{4}, {}, {3}, {}}}, {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, non_local_cols},
-          {2, {{4}, {}, {3}, {}}},
-          {4, {{}, {}, {}, {}}}}}};
-
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_mapping{
-        {false,
-         {{1, {{0}, {0, 1}, {0, 1}, {0}}},
-          {2, {{0}, {}, {0}, {}}},
-          {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, {{0}, {0, 1}, {0, 1}, {0}}},
-          {2, {{0}, {}, {0}, {}}},
-          {4, {{}, {}, {}, {}}}}}};
-};
-
-/* @brief Test fixture class for 1D mesh
- *
- * The mesh has the following structure
- *         local ids          |   global ids
- * cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
- *   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
- *        ------+------       |  --------+--------
- *        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
- *   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
- */
-class RepartitionerFixture2D
-    : public RepartitionerFixture,
-      public testing::WithParamInterface<std::tuple<bool, int>> {
-public:
-    label local_size = 4;
-
-    vec_vec idxs{
-        {0, 2, 0, 1},  // rank 0
-        {1, 3, 0, 1},  // rank 1
-        {2, 3, 0, 2},  // rank 2
-        {2, 3, 1, 3}   // rank 3
-    };
-
-    std::vector<label> rows{0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3};
-    std::vector<label> cols{0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3};
-    std::vector<label> mapping{8, 0, 1, 4, 9, 2, 5, 10, 3, 6, 7, 11};
-    std::vector<gko::span> spans{gko::span{0, 13}};
-
-    // Setup data
-    // non local data
-    vec_vec non_local_rows{
-        {1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}};
-    // non local columns in global idxs
-    vec_vec non_local_cols{
-        {4, 6, 8, 9}, {1, 3, 12, 13}, {2, 3, 12, 14}, {6, 7, 9, 11}};
-    vec_vec non_local_mapping{
-        {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}};
-    // communication partners (ranks)
-    vec_vec comm_target_ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
-
-    std::vector<std::vector<gko::span>> non_local_spans{
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 1
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 2
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 3
-    };
-
-    vec_vec non_local_ranks{
-        {1},     // rank 0
-        {0, 2},  // rank 1
-        {1, 3},  // rank 2
-        {2},     // rank 3
-    };
-
-    // expected values
-    std::map<label, vec> exp_local_nnz{
-        {1, {12, 12, 12, 12}}, {2, {28, 0, 28, 0}}, {4, {64, 0, 0, 0}}};
-    std::map<label, vec> exp_local_dim{
-        {1, {4, 4, 4, 4}}, {2, {8, 0, 8, 0}}, {4, {16, 0, 0, 0}}};
-
-    std::map<label, std::vector<label>> exp_non_local_nnz{
-        {1, {4, 4, 4, 4}}, {2, {4, 0, 4, 0}}, {4, {0, 0, 0, 0}}};
-
-    // expected local row indices in local indices
-    // first map fused true/false
-    // local row indices
-    /*
-     * The mesh has the following structure
-     *         local ids          |   global ids
-     * cells: [ 2 3 | 6 7 ]       |  [ 10 11 | 14 15 ]
-     *   2    [ 0 1 | 4 5 ]  3    |  [  8  9 | 12 13 ]
-     *        ------+------       |  --------+--------
-     *        [ 2 3 | 6 7 ]       |  [  2  3 |  6  7 ]
-     *   0    [ 0 1 | 4 5 ]  1    |  [  0  1 |  4  5 ]
-     *   */
-    vec nf_rows_2 = {0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4,
-                     4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 4, 6, 1, 3};
-    vec nf_cols_2 = {0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3, 4, 5,
-                     6, 4, 5, 7, 4, 6, 7, 5, 6, 7, 1, 3, 4, 6};
-    // [upper [0-3|12-15], lower [4,7|16-19], diag [8, 11|20-23], interfaces
-    // [24-27]]
-    vec nf_map_2 = {8,  0,  1,  4,  9,  2,
-                    5,  10, 3,  6,  7,  11,  // here first sd is done
-                    20, 12, 13, 16, 21, 14,
-                    17, 22, 15, 18, 19, 23,  // here interfaces start
-                    0,  1,  2,  3};
-    vec f_rows_2 = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3,
-                    4, 4, 4, 4, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7};
-    vec f_cols_2 = {0, 1, 2, 0, 1, 3, 4, 0, 2, 3, 1, 2, 3, 6,
-                    1, 4, 5, 6, 4, 5, 7, 3, 4, 6, 7, 5, 6, 7};
-
-    // rank 0                          | rank 1
-    // [upper, lower, diag, interfaces | upper, lower, diag, interfaces]
-    // vec f_map_2 = {8,  0,  1,  4,  9,  2,  24, 5,  10, 3,  6,  7,  11, 24,
-    //                26, 20, 12, 13, 16, 21, 14, 26, 17, 22, 15, 18, 19, 23};
-    //                the one below is broken
-    vec f_map_2 = {8,  0,  1,  4,  9,  2,  24, 5,  10, 3,  6,  7,  11, 24,
-                   26, 20, 12, 13, 16, 21, 14, 26, 17, 22, 15, 18, 19, 23};
-    /*
-     * The mesh has the following structure
-     *         local ids          |   global ids
-     * cells: [10 11|14 15]       |  [ 10 11 | 14 15 ]
-     *   2    [ 8 9 |12 13]  3    |  [  8  9 | 12 13 ]
-     *        ------+------       |  --------+--------
-     *        [ 2 3 | 6 7 ]       |  [  2  3 |  6  7 ]
-     *   0    [ 0 1 | 4 5 ]  1    |  [  0  1 |  4  5 ]
-     *   */
-    // NOTE the interfaces are in order of the comm_target indices
-    // ie [0, 0, 0 (1), (1), 2 (2), 2 (3), 2 (2), 2 (3)]
-    vec nf_rows_4 = {0,  0,  0,  1,  1,  1,  2,  2,  2,  3,  3,  3,  4,
-                     4,  4,  5,  5,  5,  6,  6,  6,  7,  7,  7,  8,  8,
-                     8,  9,  9,  9,  10, 10, 10, 11, 11, 11, 12, 12, 12,
-                     13, 13, 13, 14, 14, 14, 15, 15, 15, 4,  6,  8,  9,
-                     1,  3,  12, 13, 2,  3,  12,  14,  6,  7,  9,  11};
-    vec nf_cols_4 = {0,  1,  2,  0,  1,  3,  0,  2,  3,  1,  2,  3,  4,
-                     5,  6,  4,  5,  7,  4,  6,  7,  5,  6,  7,  8,  9,
-                     10, 8,  9,  11, 8,  10, 11, 9,  10, 11, 12, 13, 14,
-                     12, 13, 15, 12, 14, 15, 13, 14, 15, 1,  3,  2,  3,
-                     4,  6,  6,  7,  8,  9,  9,  11, 12, 13, 12, 14};
-    vec nf_map_4{8,  0,  1,  4,  9,  2,  5,  10, 3,  6,  7,  11,  // 1st sd done
-                 20, 12, 13, 16, 21, 14, 17, 22, 15, 18, 19, 23,  // 2nd sd done
-                 32, 24, 25, 28, 33, 26, 29, 34, 27, 30, 31, 35,  // 3rd sd done
-                 44, 36, 37, 40, 45, 38, 41, 46, 39, 42, 43, 47,  // 4th sd done
-                 0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11,
-                 12, 13, 14, 15};
-    // fused
-    vec f_rows_4 = {0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,
-                    3,  3,  3,  4,  4,  4,  4,  5,  5,  5,  6,  6,  6,
-                    6,  6,  7,  7,  7,  7,  8,  8,  8,  8,  9,  9,  9,
-                    9,  9,  10, 10, 10, 11, 11, 11, 11, 12, 12, 12, 12,
-                    12, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15};
-    vec f_cols_4 = {0,  1,  2,  0,  1,  3,  4,  0,  2,  3,  8,  1,  2,
-                    3,  6,  9,  1,  4,  5,  6,  4,  5,  7,  3,  4,  6,
-                    7,  12, 5,  6,  7,  13, 2,  8,  9,  10, 3,  8,  9,
-                    11, 12, 8,  10, 11, 9,  10, 11, 14, 6,  9,  12, 13,
-                    14, 7,  12, 13, 15, 11, 12, 14, 15, 13, 14, 15};
-    vec f_map_4{0,  1, 2, 3,  16, 17, 4,  5,  6,  7,  18,
-                19, 8, 9, 10, 11, 20, 21, 12, 13, 14, 15};
-
-    std::map<bool, std::map<label, vec_vec>> exp_local_rows{
-        {false,
-         {{1, {rows, rows, rows, rows}},
-          {2, {nf_rows_2, {}, nf_rows_2, {}}},
-          {4, {nf_rows_4, {}, {}, {}}}}},
-        {true,
-         {{1, {rows, rows, rows, rows}},
-          {2, {f_rows_2, {}, f_rows_2, {}}},
-          {4, {f_rows_4, {}, {}, {}}}}}};
-
-    std::map<bool, std::map<label, vec_vec>> exp_local_cols{
-        {false,
-         {{1, {cols, cols, cols, cols}},
-          {2, {nf_cols_2, {}, nf_cols_2, {}}},
-          {4, {nf_cols_4, {}, {}, {}}}}},
-        {true,
-         {{1, {cols, cols, cols, cols}},
-          {2, {f_cols_2, {}, f_cols_2, {}}},
-          {4, {f_cols_4, {}, {}, {}}}}}};
-
-    std::map<bool, std::map<label, vec_vec>> exp_local_mapping{
-        {false,
-         {{1, {mapping, mapping, mapping, mapping}},
-          {2, {nf_map_2, {}, nf_map_2, {}}},
-          {4, {nf_map_4, {}, {}, {}}}}},
-        {true,
-         {{1, {mapping, mapping, mapping, mapping}},
-          {2, {f_map_2, {}, f_map_2, {}}},
-          {4, {f_map_4, {}, {}, {}}}}}};
-
-    // [start,end) indices of interface/submatrices
-    std::map<bool, std::map<label, vec_vec>> exp_local_spans_begin{
-        {false,
-         {{1, {{0}, {0}, {0}, {0}}},
-          {2, {{0, 24, 26}, {}, {0, 24, 26}, {}}},
-          {4, {{0, 48, 50, 52, 54, 56, 58, 60, 62}, {}, {}, {}}}}},
-        {true,
-         {{1, {{0}, {0}, {0}, {0}}},
-          {2, {{0}, {}, {0}, {}}},
-          {4, {{0}, {}, {}, {}}}}}};
-
-    // [start,end) indices of interface/submatrices
-    std::map<bool, std::map<label, vec_vec>> exp_local_spans_end{
-        {false,
-         {{1, {{13}, {13}, {13}, {13}}},
-          {2, {{24, 26, 28}, {}, {24, 26, 28}, {}}},
-          {4, {{48, 50, 52, 54, 56, 58, 60, 62, 64}, {}, {}, {}}}}},
-        {true,
-         {{1, {{13}, {13}, {13}, {13}}},
-          {2, {{10}, {}, {10}, {}}},
-          {4, {{22}, {}, {}, {}}}}}};
-
-    // non local data
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_rows{
-        {false,
-         {{1, non_local_rows},
-          {2, {{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}}},
-          {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, non_local_rows},
-          {2, {{}, {}, {}, {}}},  // missing
-          {4, {{}, {}, {}, {}}}}}};
-
-    // non local cols are in global indices
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_cols{
-        {false,
-         {{1, non_local_cols},
-          {2, {{8, 9, 12, 13}, {}, {2, 3, 6, 7}, {}}},
-          {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, non_local_cols}, {2, {{}, {}, {}, {}}}, {4, {{}, {}, {}, {}}}}}};
-
-    std::map<bool, std::map<label, vec_vec>> exp_non_local_mapping{
-        {false,
-         {{1, {{0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}}},
-          {2, {{0, 1, 2, 3}, {}, {0, 1, 2, 3}, {}}},
-          {4, {{}, {}, {}, {}}}}},
-        {true,
-         {{1, non_local_cols}, {2, {{}, {}, {}, {}}}, {4, {{}, {}, {}, {}}}}}};
-};
-
-INSTANTIATE_TEST_SUITE_P(RepartitionerFixture1DInstantiation,
-                         RepartitionerFixture1D,
-                         testing::Combine(testing::Values(false),
-                                          testing::Values(1, 2, 4)));
-INSTANTIATE_TEST_SUITE_P(
-    RepartitionerFixture2DInstantiation, RepartitionerFixture2D,
-    testing::Combine(testing::Values(false), testing::Values(1, 2, 4))  //,
-    // [](const testing::TestParamInfo<MyTestSuite::ParamType> &info) {
-    //     // Can use info.param here to generate the test suffix
-    //     std::string name = "foo";
-    //     return name;
-    // }
-);
-
-TEST_P(RepartitionerFixture1D, can_create_repartitioner)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    // Act
-    auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, exec, fused);
-    // Assert
-    EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
-}
-
-TEST_P(RepartitionerFixture1D, has_correct_properties_for_n_rank)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    // Act
-    auto repartitioner =
-        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
-
-    // Assert
-    EXPECT_EQ(repartitioner.is_owner(exec),
-              (rank % ranks_per_gpu == 0) ? true : false);
-
-    EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, exec),
-        (repartitioner.is_owner(exec) ? ranks_per_gpu * local_size : 0));
-}
-
-TEST_P(RepartitionerFixture2D, can_convert_to_global)
-{
-    // Arrange
-    using namespace gko::experimental::distributed;
-    auto ref_exec = exec.get_ref_exec();
-    auto partition = gko::share(
-        build_partition_from_local_size<label, label>(ref_exec, comm, 4));
-
-    vec_vec exp_global_idxs{
-        {4, 6, 8, 9},    // rank 0
-        {1, 3, 12, 13},  // rank 1
-        {2, 3, 12, 14},  // rank 2
-        {6, 7, 9, 11}    // rank 3
-    };
-
-    // Act
-    auto global_idxs = detail::convert_to_global(
-        partition, idxs[rank].data(), non_local_spans[rank], comm_target_ids[rank]);
-
-    // Assert
-    EXPECT_EQ(global_idxs, exp_global_idxs[rank]);
-}
-
-
-TEST_P(RepartitionerFixture1D, can_exchange_spans_and_ranks_for_n_ranks)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-
-    std::vector<gko::span> spans{{0, 5}, {5, 10}};
-    std::vector<label> ranks{rank + 1, rank + 2};
-
-    std::map<label, vec_vec> exp_ranks;
-    exp_ranks.emplace(1, vec_vec{{1, 2}, {2, 3}, {3, 4}, {4, 5}});
-    exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
-    exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
-
-    std::map<label, vec_vec> exp_spans_begin;
-    exp_spans_begin.emplace(1, vec_vec{{0, 5}, {0, 5}, {0, 5}, {0, 5}});
-    exp_spans_begin.emplace(2, vec_vec{{0, 5, 10, 15}, {}, {0, 5, 10, 15}, {}});
-    exp_spans_begin.emplace(
-        4, vec_vec{{0, 5, 10, 15, 20, 25, 30, 35}, {}, {}, {}});
-    std::map<label, vec_vec> exp_spans_end;
-    exp_spans_end.emplace(1, vec_vec{{5, 10}, {5, 10}, {5, 10}, {5, 10}});
-    exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 20}, {}});
-    exp_spans_end.emplace(4,
-                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
-
-    // Act
-    auto [new_spans, origins, gathered_ranks] =
-        detail::exchange_spans_ranks(exec, ranks_per_gpu, spans, ranks);
-
-    std::vector<label> res_spans_begin{};
-    std::vector<label> res_spans_end{};
-
-    for (auto span : new_spans) {
-        res_spans_begin.push_back(span.begin);
-        res_spans_end.push_back(span.end);
-    }
-
-    // Assert
-    ASSERT_EQ(res_spans_begin, exp_spans_begin[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_spans_end, exp_spans_end[ranks_per_gpu][rank]);
-}
-
-TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    auto repartitioner =
-        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
-    auto ref_exec = exec.get_ref_exec();
-
-    std::vector<label> ranks{rank};
-    auto local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans);
-
-    auto non_local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{2, non_local_ranks[rank].size()},
-        non_local_rows[rank], non_local_cols[rank], non_local_mapping[rank],
-        non_local_spans[rank]);
-
-    // Act
-    auto [repart_local, repart_non_local, tracking] =
-        repartitioner.repartition_sparsity(exec, local_sparsity,
-                                           non_local_sparsity,
-                                           comm_target_ids[rank], fused);
-    // Assert
-    // local properties
-    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[1], exp_local_dim[ranks_per_gpu][rank]);
-
-    auto res_local_rows = convert_to_vector(repart_local->row_idxs);
-    auto res_local_cols = convert_to_vector(repart_local->col_idxs);
-    auto res_local_mapping = convert_to_vector(repart_local->ldu_mapping);
-    ASSERT_EQ(res_local_rows, exp_local_rows[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_cols, exp_local_cols[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_mapping, exp_local_mapping[fused][ranks_per_gpu][rank]);
-
-    std::vector<label> res_local_spans_begin{};
-    std::vector<label> res_local_spans_end{};
-    for (auto [begin, end] : repart_local->spans) {
-        res_local_spans_begin.push_back(begin);
-        res_local_spans_end.push_back(end);
-    }
-    ASSERT_EQ(res_local_spans_begin,
-              exp_local_spans_begin[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_spans_end,
-              exp_local_spans_end[fused][ranks_per_gpu][rank]);
-
-    // non local properties
-    ASSERT_EQ(repart_non_local->num_nnz,
-              exp_non_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
-
-    auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
-    auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);
-    ASSERT_EQ(res_non_local_rows,
-              exp_non_local_rows[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_non_local_cols,
-              exp_non_local_cols[fused][ranks_per_gpu][rank]);
-    auto res_non_local_mapping =
-        convert_to_vector(repart_non_local->ldu_mapping);
-    ASSERT_EQ(res_non_local_mapping,
-              exp_non_local_mapping[fused][ranks_per_gpu][rank]);
-}
-
-TEST_P(RepartitionerFixture2D, can_repartition_sparsity_pattern_2D_for_n_ranks)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    auto repartitioner =
-        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
-    auto ref_exec = exec.get_ref_exec();
-
-    std::vector<label> ranks{rank};
-    auto local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{4, 4}, rows, cols, mapping, spans);
-
-    auto non_local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{4, 4}, non_local_rows[rank], non_local_cols[rank],
-        non_local_mapping[rank], non_local_spans[rank]);
-
-    // Act
-    auto [repart_local, repart_non_local, tracking] =
-        repartitioner.repartition_sparsity(exec, local_sparsity,
-                                           non_local_sparsity,
-                                           comm_target_ids[rank], fused);
-    // Assert
-    // local properties
-    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[1], exp_local_dim[ranks_per_gpu][rank]);
-
-    auto res_local_rows = convert_to_vector(repart_local->row_idxs);
-    auto res_local_cols = convert_to_vector(repart_local->col_idxs);
-    auto res_local_mapping = convert_to_vector(repart_local->ldu_mapping);
-    ASSERT_EQ(res_local_rows.size(),
-              exp_local_rows[fused][ranks_per_gpu][rank].size());
-    for (size_t i = 0; i < res_local_rows.size(); i++) {
-        ASSERT_EQ(res_local_rows[i],
-                  exp_local_rows[fused][ranks_per_gpu][rank][i])
-            << " failed at index " << i;
-    }
-    ASSERT_EQ(res_local_cols.size(),
-              exp_local_cols[fused][ranks_per_gpu][rank].size());
-    for (size_t i = 0; i < res_local_cols.size(); i++) {
-        ASSERT_EQ(res_local_cols[i],
-                  exp_local_cols[fused][ranks_per_gpu][rank][i])
-            << " failed at index " << i;
-    }
-    ASSERT_EQ(res_local_mapping.size(),
-              exp_local_cols[fused][ranks_per_gpu][rank].size());
-    for (size_t i = 0; i < res_local_cols.size(); i++) {
-        ASSERT_EQ(res_local_cols[i],
-                  exp_local_cols[fused][ranks_per_gpu][rank][i])
-            << " failed at index " << i;
-    }
-
-    std::vector<label> res_local_spans_begin{};
-    std::vector<label> res_local_spans_end{};
-    for (auto [begin, end] : repart_local->spans) {
-        res_local_spans_begin.push_back(begin);
-        res_local_spans_end.push_back(end);
-    }
-    ASSERT_EQ(res_local_spans_begin,
-              exp_local_spans_begin[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_spans_end,
-              exp_local_spans_end[fused][ranks_per_gpu][rank]);
-
-    // non local properties
-    ASSERT_EQ(repart_non_local->num_nnz,
-              exp_non_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
-
-    auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
-    auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);
-    ASSERT_EQ(res_non_local_rows,
-              exp_non_local_rows[fused][ranks_per_gpu][rank]);
-    ASSERT_EQ(res_non_local_cols,
-              exp_non_local_cols[fused][ranks_per_gpu][rank]);
-    auto res_non_local_mapping =
-        convert_to_vector(repart_non_local->ldu_mapping);
-    ASSERT_EQ(res_non_local_mapping,
-              exp_non_local_mapping[fused][ranks_per_gpu][rank]);
-}
-
-
-TEST_P(RepartitionerFixture1D, can_repartition_1D_comm_pattern_for_n_ranks)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    auto repartitioner =
-        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
-    auto ref_exec = exec.get_ref_exec();
-
-    // expected communication ranks
-    std::map<label, vec_vec> exp_res_ids{};
-
-    exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
-    // only communication partners are 0-2 and 2-0
-    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
-    // no communication if all ranks are repartitioned to single owner
-    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
-
-    // expected communication sizes
-    std::map<label, vec_vec> exp_res_sizes;
-    exp_res_sizes.emplace(1, vec_vec({{1}, {1, 1}, {1, 1}, {1}}));
-    exp_res_sizes.emplace(2, vec_vec({{1}, {}, {1}, {}}));
-    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
-
-    // expected rows
-    vec_vec_vec rows{{{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
-    std::map<label, vec_vec> exp_res_rows;
-    exp_res_rows.emplace(1, vec_vec{{1}, {0, 1}, {0, 1}, {0}});
-    // after repartitioning with 2 ranks_per_gpu [ 0 1 2 3 | 0 1 2 3 ] <-
-    // local row ids
-    exp_res_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
-    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <-
-    // local row ids
-    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
-
-    // the original comm_pattern
-    auto comm_pattern =
-        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
-
-    // Act
-    auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(exec, comm_pattern);
-
-    // Assert
-    auto res_ids = repart_comm_pattern->target_ids;
-    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
-
-    auto res_sizes = repart_comm_pattern->target_sizes;
-    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
-
-    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
-    auto res_rows = total_rank_send_idx;
-    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
-}
-
-TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
-{
-    // Arrange
-    auto [fused, ranks_per_gpu] = GetParam();
-    auto repartitioner =
-        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
-    auto ref_exec = exec.get_ref_exec();
-
-    // expected communication ranks
-    std::map<label, vec_vec> exp_res_ids{};
-    exp_res_ids[1] = comm_target_ids;  // in the ranks_per_gpu==1 case nothing changes
-    // only communication partners are 0-2 and 2-0
-    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
-    // no communication if all ranks are repartitioned to single owner
-    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
-
-    // expected communication sizes
-    std::map<label, vec_vec> exp_res_sizes;
-    exp_res_sizes.emplace(1, vec_vec({{2, 2}, {2, 2}, {2, 2}, {2, 2}}));
-    exp_res_sizes.emplace(2, vec_vec({{4}, {}, {4}, {}}));
-    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
-
-    // expected rows
-    vec_vec_vec rows{
-        {{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
-
-    std::map<label, vec_vec> exp_res_rows;
-    exp_res_rows.emplace(
-        1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}});
-    // after repartitioning with 2 ranks_per_gpu
-    // cells: [ 2 3  6 7 ]
-    //   2    [ 0 1  4 5 ]  3
-    //        -------------
-    //        [ 2 3  6 7 ]
-    //   0    [ 0 1  4 5 ]  1
-    exp_res_rows.emplace(2, vec_vec{{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}});
-    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <-
-    // local row ids
-    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
-
-    std::map<label, vec_vec> exp_gather_idx;
-    exp_gather_idx.emplace(
-        1, vec_vec{{0, 2, 0, 1}, {1, 3, 0, 1}, {2, 3, 0, 2}, {2, 3, 1, 3}});
-    exp_gather_idx.emplace(2, vec_vec{{0, 1, 4, 5}, {}, {2, 3, 6, 7}, {}});
-    exp_gather_idx.emplace(4, vec_vec{{}, {}, {}, {}});
-
-    // the original comm_pattern
-    auto comm_pattern =
-        std::make_shared<CommunicationPattern>(exec, comm_target_ids[rank], rows[rank]);
-
-    // Act
-    auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(exec, comm_pattern);
-
-    // Assert
-    auto res_ids = repart_comm_pattern->target_ids;
-
-    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
-
-    auto res_sizes = repart_comm_pattern->target_sizes;
-    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
-
-    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
-    auto res_rows = total_rank_send_idx;
-    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
-
-    auto total_recv_gather_idx =
-        repart_comm_pattern->compute_recv_gather_idxs(exec);
-    auto res_gather_idx = convert_to_vector(total_recv_gather_idx);
-    EXPECT_EQ(res_gather_idx, exp_gather_idx[ranks_per_gpu][rank]);
-}
-
-
-int main(int argc, char *argv[])
-{
-    int result = 0;
-
-    ::testing::InitGoogleTest(&argc, argv);
-
-    my_argc = argc;
-    my_argv = argv;
-
-    result = RUN_ALL_TESTS();
-
-    return result;
-}
diff --git a/unitTests/Repartitioner.H b/unitTests/Repartitioner.H
new file mode 100644
index 000000000..53ae19a52
--- /dev/null
+++ b/unitTests/Repartitioner.H
@@ -0,0 +1,86 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "OGL/Repartitioner.H"
+#include "OGL/common.H"
+
+#include "gtest/gtest.h"
+
+#include "mpi.h"
+
+#include <cstdlib>
+
+//---------------------------------------------
+// some_header.h
+extern int my_argc;
+extern char **my_argv;
+// eof
+//---------------------------------------------
+
+//---------------------------------------------
+// main.cpp
+int my_argc;
+char **my_argv;
+
+using vec = std::vector<label>;
+using vec_vec = std::vector<std::vector<label>>;
+using vec_vec_vec = std::vector<vec_vec>;
+
+class RepartitionerEnvironment : public testing::Environment {
+public:
+    void SetUp()
+    {
+        args_ = std::make_shared<Foam::argList>(my_argc, my_argv);
+        time = std::make_shared<Foam::Time>("controlDict", *args_.get());
+        db = std::make_shared<Foam::objectRegistry>(*time.get());
+        Foam::dictionary dict;
+        dict.add("executor", "reference");
+        exec = std::make_shared<ExecutorHandler>(time->thisDb(), dict, "dummy",
+                                                 true);
+
+        if (exec->get_communicator()->size() < 2) {
+            std::cout << "At least 2 CPU processes should be used!"
+                      << std::endl;
+            std::abort();
+        }
+        // delete listener on ranks != 0
+        // to clean up output
+        ::testing::TestEventListeners &listeners =
+            ::testing::UnitTest::GetInstance()->listeners();
+        // if (exec->get_communicator()->rank() != 0) {
+        //     delete listeners.Release(listeners.default_result_printer());
+        // }
+    }
+
+    std::shared_ptr<Foam::argList> args_;
+    std::shared_ptr<const Foam::Time> time;
+    std::shared_ptr<Foam::objectRegistry> db;
+    std::shared_ptr<ExecutorHandler> exec;
+};
+
+const testing::Environment *global_env =
+    AddGlobalTestEnvironment(new RepartitionerEnvironment);
+
+class RepartitionerFixture : public testing::Test {
+public:
+    ExecutorHandler exec =
+        *((RepartitionerEnvironment *)global_env)->exec.get();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(exec.get_communicator().get());
+};
+
+int main(int argc, char *argv[])
+{
+    int result = 0;
+
+    ::testing::InitGoogleTest(&argc, argv);
+
+    my_argc = argc;
+    my_argv = argv;
+
+    result = RUN_ALL_TESTS();
+
+    return result;
+}
diff --git a/unitTests/Repartitioner1D.C b/unitTests/Repartitioner1D.C
new file mode 100644
index 000000000..cb1ada505
--- /dev/null
+++ b/unitTests/Repartitioner1D.C
@@ -0,0 +1,388 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "Repartitioner.H"
+
+/* @brief Test fixture class for 1D mesh
+ *
+ * The mesh has the following structure
+ * ranks:    0     1     2     3
+ * cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+ * cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
+ */
+class RepartitionerFixture1D
+    : public RepartitionerFixture,
+      public testing::WithParamInterface<std::tuple<bool, int>> {
+public:
+    label local_size = 2;  // matrix rows, same for all ranks
+
+    // local data
+    std::vector<label> rows{0, 0, 1, 1};
+    std::vector<label> cols{0, 1, 0, 1};
+    // mapping is [u, l, d]
+    std::vector<label> mapping{2, 0, 1, 3};
+    std::vector<gko::span> spans{gko::span{0, 5}};
+
+    // Setup data
+    // non local data
+    vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
+    // non local columns in global idxs
+    vec_vec non_local_cols{{2}, {1, 4}, {3, 6}, {5}};
+    vec_vec non_local_mapping{{0}, {0, 1}, {0, 1}, {0}};
+    // communication partners (ranks)
+    vec_vec ids{{1}, {0, 2}, {1, 3}, {2}};
+
+    std::vector<std::vector<gko::span>> non_local_spans{
+        {gko::span{0, 1}},                   // rank 0
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 1
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 2
+        {gko::span{0, 1}},                   // rank 3
+    };
+
+    vec_vec non_local_ranks{
+        {1},     // rank 0
+        {0, 2},  // rank 1
+        {1, 3},  // rank 2
+        {2},     // rank 3
+    };
+
+    // Expected results
+    // number of non-zeros of each sparsity pattern
+    std::map<label, std::vector<label>> exp_local_nnz{
+        {1, {4, 4, 4, 4}},
+        {2, {10, 0, 10, 0}},
+        {4, {22, 0, 0, 0}},
+    };
+
+    std::map<label, std::vector<label>> exp_non_local_nnz{
+        {1, {1, 2, 2, 1}},
+        {2, {1, 0, 1, 0}},
+        {4, {0, 0, 0, 0}},
+    };
+
+    std::vector<std::vector<label>> comm_target_ids{{1}, {0, 2}, {1, 3}, {2}};
+    // expected local row indices in local indices
+    // first map fused true/false
+    // local row indices
+    // repartitioned 2 [ 0 1 , 2 3 | 0 1 , 2 3 ] | new  boundary , old boundary
+    // repartitioned 4 [ 0 1 , 2 3 , 4 5 , 6 7 ] | new  boundary , old boundary
+    // the last two elements are now local interfaces, they are in in order
+    // of target_ids
+    vec nf_rows_2 = {0, 0, 1, 1, 2, 2, 3, 3, 2, 1};
+    vec nf_rows_4 = {0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5,
+                     5, 6, 6, 7, 7, 2, 1, 4, 3, 6, 5};
+    // fused
+    vec f_rows_2 = {0, 0, 1, 1, 1, 2, 2, 2, 3, 3};
+    vec f_rows_4 = {0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3,
+                    4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7};
+    std::map<bool, std::map<label, vec_vec>> exp_local_rows{
+        {false,
+         {{1, {rows, rows, rows, rows}},
+          {2, {nf_rows_2, {}, nf_rows_2, {}}},
+          {4, {nf_rows_4, {}, {}, {}}}}},
+        {true,
+         {{1, {rows, rows, rows, rows}},
+          {2, {f_rows_2, {}, f_rows_2, {}}},
+          {4, {f_rows_4, {}, {}, {}}}}}};
+
+    vec nf_cols_2 = {0, 1, 0, 1, 2, 3, 2, 3, 1, 2};
+    vec nf_cols_4 = {0, 1, 0, 1, 2, 3, 2, 3, 4, 5, 4,
+                     5, 6, 7, 6, 7, 1, 2, 3, 4, 5, 6};
+    // fused
+    vec f_cols_2 = {0, 1, 0, 1, 2, 1, 2, 3, 2, 3};
+    vec f_cols_4 = {0, 1, 0, 1, 2, 1, 2, 3, 2, 3, 4,
+                    3, 4, 5, 4, 5, 6, 5, 6, 7, 6, 7};
+    // first map fused true/false
+    // expected local col indices in local indices
+    std::map<bool, std::map<label, vec_vec>> exp_local_cols{
+        {false,
+         {{1, {cols, cols, cols, cols}},
+          {2, {nf_cols_2, {}, nf_cols_2, {}}},
+          {4, {nf_cols_4, {}, {}, {}}}}},
+        {true,
+         {{1, {cols, cols, cols, cols}},
+          {2, {f_cols_2, {}, f_cols_2, {}}},
+          {4, {f_cols_4, {}, {}, {}}}}}};
+
+    vec nf_map_2{2, 0, 1, 3, 6, 4, 5, 7,  // here the interface values start
+                 8, 9};                   // <- they are currently unused
+    vec nf_map_4{2,  0,  1,  3,  6,  4,  5,  7,
+                 10, 8,  9,  11, 14, 12, 13, 15,  // here the interface start
+                 16, 17, 18, 19, 20, 21};         // <- unused values because we
+                                                  // dont permute atm
+    // [u l d | u l d ], [ i | i ]
+    vec f_map_2{2, 0, 1, 3, 9, 8, 6, 4, 5, 7};
+    vec f_map_4{2,  0,  1, 3, 17, 16, 6,  4,  5,  7,  19,
+                18, 10, 8, 9, 11, 21, 20, 14, 12, 13, 15};
+    std::map<bool, std::map<label, vec_vec>> exp_local_mapping{
+        {false,
+         {{1, {mapping, mapping, mapping, mapping}},
+          {2, {nf_map_2, {}, nf_map_2, {}}},
+          {4, {nf_map_4, {}, {}, {}}}}},
+        {true,
+         {{1, {mapping, mapping, mapping, mapping}},
+          {2, {f_map_2, {}, f_map_2, {}}},
+          {4, {f_map_4, {}, {}, {}}}}}};
+
+    // [start,end) indices of interface/submatrices
+    std::map<bool, std::map<label, vec_vec>> exp_local_spans_begin{
+        {false,
+         {{1, {{0}, {0}, {0}, {0}}},
+          {2, {{0, 8, 9}, {}, {0, 8, 9}, {}}},
+          {4, {{0, 16, 17, 18, 19, 20, 21}, {}, {}, {}}}}},
+        {true,
+         {{1, {{0}, {0}, {0}, {0}}},
+          {2, {{0}, {}, {0}, {}}},
+          {4, {{0}, {}, {}, {}}}}}};
+
+    // [start,end) indices of interface/submatrices
+    std::map<bool, std::map<label, vec_vec>> exp_local_spans_end{
+        {false,
+         {{1, {{5}, {5}, {5}, {5}}},
+          {2, {{8, 9, 10}, {}, {8, 9, 10}, {}}},
+          {4, {{16, 17, 18, 19, 20, 21, 22}, {}, {}, {}}}}},
+        {true,
+         {{1, {{5}, {5}, {5}, {5}}},
+          {2, {{10}, {}, {10}, {}}},
+          {4, {{22}, {}, {}, {}}}}}};
+
+    // number of rows of local matrix
+    std::map<label, vec> exp_local_dim{
+        {1, {local_size, local_size, local_size, local_size}},
+        {2, {2 * local_size, 0, 2 * local_size, 0}},
+        {4, {4 * local_size, 0, 0, 0}},
+    };
+
+    // non local data
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_rows{
+        {false,
+         {{1, non_local_rows}, {2, {{3}, {}, {0}, {}}}, {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, non_local_rows},
+          {2, {{3}, {}, {0}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    // non local cols are in global indices
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_cols{
+        {false,
+         {{1, non_local_cols}, {2, {{4}, {}, {3}, {}}}, {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, non_local_cols},
+          {2, {{4}, {}, {3}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_mapping{
+        {false,
+         {{1, {{0}, {0, 1}, {0, 1}, {0}}},
+          {2, {{0}, {}, {0}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, {{0}, {0, 1}, {0, 1}, {0}}},
+          {2, {{0}, {}, {0}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+};
+
+
+INSTANTIATE_TEST_SUITE_P(RepartitionerFixture1DInstantiation,
+                         RepartitionerFixture1D,
+                         testing::Combine(testing::Values(false, true),
+                                          testing::Values(1, 2, 4)),
+                         [](const auto &info) {
+                             // Can use info.param here to generate the test
+                             // suffix
+                             std::vector<std::string> names;
+                             names.emplace_back("fuse");
+                             names.emplace_back("ranks");
+                             std::string name = "fused_";
+                             name += std::to_string(std::get<0>(info.param));
+                             name += "_ranks_";
+                             name += std::to_string(std::get<1>(info.param));
+                             return name;
+                         });
+
+TEST_P(RepartitionerFixture1D, can_create_repartitioner)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+    // Act
+    auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, exec, fused);
+    // Assert
+    EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
+}
+
+TEST_P(RepartitionerFixture1D, has_correct_properties_for_n_rank)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+    // Act
+    auto repartitioner =
+        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
+
+    // Assert
+    EXPECT_EQ(repartitioner.is_owner(exec),
+              (rank % ranks_per_gpu == 0) ? true : false);
+
+    EXPECT_EQ(
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, exec),
+        (repartitioner.is_owner(exec) ? ranks_per_gpu * local_size : 0));
+}
+
+TEST_P(RepartitionerFixture1D, can_exchange_spans_and_ranks_for_n_ranks)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+
+    std::vector<gko::span> spans{{0, 5}, {5, 10}};
+    std::vector<label> ranks{rank + 1, rank + 2};
+
+    std::map<label, vec_vec> exp_ranks;
+    exp_ranks.emplace(1, vec_vec{{1, 2}, {2, 3}, {3, 4}, {4, 5}});
+    exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
+    exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
+
+    std::map<label, vec_vec> exp_spans_begin;
+    exp_spans_begin.emplace(1, vec_vec{{0, 5}, {0, 5}, {0, 5}, {0, 5}});
+    exp_spans_begin.emplace(2, vec_vec{{0, 5, 10, 15}, {}, {0, 5, 10, 15}, {}});
+    exp_spans_begin.emplace(
+        4, vec_vec{{0, 5, 10, 15, 20, 25, 30, 35}, {}, {}, {}});
+    std::map<label, vec_vec> exp_spans_end;
+    exp_spans_end.emplace(1, vec_vec{{5, 10}, {5, 10}, {5, 10}, {5, 10}});
+    exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 20}, {}});
+    exp_spans_end.emplace(4,
+                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
+
+    // Act
+    auto [new_spans, origins, gathered_ranks] =
+        detail::exchange_spans_ranks(exec, ranks_per_gpu, spans, ranks);
+
+    std::vector<label> res_spans_begin{};
+    std::vector<label> res_spans_end{};
+
+    for (auto span : new_spans) {
+        res_spans_begin.push_back(span.begin);
+        res_spans_end.push_back(span.end);
+    }
+
+    // Assert
+    ASSERT_EQ(res_spans_begin, exp_spans_begin[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_spans_end, exp_spans_end[ranks_per_gpu][rank]);
+}
+
+TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+    auto repartitioner =
+        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
+    auto ref_exec = exec.get_ref_exec();
+
+    std::vector<label> ranks{rank};
+    auto local_sparsity = std::make_shared<SparsityPattern>(
+        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans);
+
+    auto non_local_sparsity = std::make_shared<SparsityPattern>(
+        ref_exec, gko::dim<2>{2, non_local_ranks[rank].size()},
+        non_local_rows[rank], non_local_cols[rank], non_local_mapping[rank],
+        non_local_spans[rank]);
+
+    // Act
+    auto [repart_local, repart_non_local, tracking] =
+        repartitioner.repartition_sparsity(exec, local_sparsity,
+                                           non_local_sparsity,
+                                           comm_target_ids[rank], fused);
+    // Assert
+    // local properties
+    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[1], exp_local_dim[ranks_per_gpu][rank]);
+
+    auto res_local_rows = convert_to_vector(repart_local->row_idxs);
+    auto res_local_cols = convert_to_vector(repart_local->col_idxs);
+    auto res_local_mapping = convert_to_vector(repart_local->ldu_mapping);
+    ASSERT_EQ(res_local_rows, exp_local_rows[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_cols, exp_local_cols[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_mapping, exp_local_mapping[fused][ranks_per_gpu][rank]);
+
+    std::vector<label> res_local_spans_begin{};
+    std::vector<label> res_local_spans_end{};
+    for (auto [begin, end] : repart_local->spans) {
+        res_local_spans_begin.push_back(begin);
+        res_local_spans_end.push_back(end);
+    }
+    ASSERT_EQ(res_local_spans_begin,
+              exp_local_spans_begin[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_spans_end,
+              exp_local_spans_end[fused][ranks_per_gpu][rank]);
+
+    // non local properties
+    ASSERT_EQ(repart_non_local->num_nnz,
+              exp_non_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
+
+    auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
+    auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);
+    ASSERT_EQ(res_non_local_rows,
+              exp_non_local_rows[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_cols,
+              exp_non_local_cols[fused][ranks_per_gpu][rank]);
+    auto res_non_local_mapping =
+        convert_to_vector(repart_non_local->ldu_mapping);
+    ASSERT_EQ(res_non_local_mapping,
+              exp_non_local_mapping[fused][ranks_per_gpu][rank]);
+}
+
+TEST_P(RepartitionerFixture1D, can_repartition_comm_pattern)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+    auto repartitioner =
+        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
+    auto ref_exec = exec.get_ref_exec();
+
+    // expected communication ranks
+    std::map<label, vec_vec> exp_res_ids{};
+
+    exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
+    // only communication partners are 0-2 and 2-0
+    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
+    // no communication if all ranks are repartitioned to single owner
+    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // expected communication sizes
+    std::map<label, vec_vec> exp_res_sizes;
+    exp_res_sizes.emplace(1, vec_vec({{1}, {1, 1}, {1, 1}, {1}}));
+    exp_res_sizes.emplace(2, vec_vec({{1}, {}, {1}, {}}));
+    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
+
+    // expected rows
+    vec_vec_vec rows{{{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+    std::map<label, vec_vec> exp_res_rows;
+    exp_res_rows.emplace(1, vec_vec{{1}, {0, 1}, {0, 1}, {0}});
+    // after repartitioning with 2 ranks_per_gpu [ 0 1 2 3 | 0 1 2 3 ] <-
+    // local row ids
+    exp_res_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <-
+    // local row ids
+    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // the original comm_pattern
+    auto comm_pattern =
+        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
+
+    // Assert
+    auto res_ids = repart_comm_pattern->target_ids;
+    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
+
+    auto res_sizes = repart_comm_pattern->target_sizes;
+    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+
+    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    auto res_rows = total_rank_send_idx;
+    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
+}
diff --git a/unitTests/Repartitioner2D.C b/unitTests/Repartitioner2D.C
index 30cec7f51..193a670ed 100644
--- a/unitTests/Repartitioner2D.C
+++ b/unitTests/Repartitioner2D.C
@@ -2,39 +2,400 @@
 //
 // SPDX-License-Identifier: GPL-3.0-or-later
 
-TEST_P(RepartitionerFixture2D, can_repartition_sparsity_pattern_2D_for_n_ranks_unfused)
+#include "Repartitioner.H"
+
+/* @brief Test fixture class for 2D mesh
+ *
+ * The mesh has the following structure
+ *         local ids          |   global ids
+ * cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
+ *   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
+ *        ------+------       |  --------+--------
+ *        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
+ *   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
+ */
+class RepartitionerFixture2D
+    : public RepartitionerFixture,
+      public testing::WithParamInterface<std::tuple<bool, int>> {
+public:
+    label local_size = 4;
+
+    vec_vec idxs{
+        {0, 2, 0, 1},  // rank 0
+        {1, 3, 0, 1},  // rank 1
+        {2, 3, 0, 2},  // rank 2
+        {2, 3, 1, 3}   // rank 3
+    };
+
+    std::vector<label> rows{0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3};
+    std::vector<label> cols{0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3};
+    std::vector<label> mapping{8, 0, 1, 4, 9, 2, 5, 10, 3, 6, 7, 11};
+    std::vector<gko::span> spans{gko::span{0, 13}};
+
+    // Setup data
+    // non local data
+    vec_vec non_local_rows{
+        {1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}};
+    // non local columns in global idxs
+    vec_vec non_local_cols{
+        {4, 6, 8, 9}, {1, 3, 12, 13}, {2, 3, 12, 14}, {6, 7, 9, 11}};
+    vec_vec non_local_mapping{
+        {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}};
+    // communication partners (ranks)
+    vec_vec comm_target_ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
+
+    std::vector<std::vector<gko::span>> non_local_spans{
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 1
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 2
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 3
+    };
+
+    vec_vec non_local_ranks{
+        {1},     // rank 0
+        {0, 2},  // rank 1
+        {1, 3},  // rank 2
+        {2},     // rank 3
+    };
+
+    // expected values
+    std::map<label, vec> exp_local_nnz{
+        {1, {12, 12, 12, 12}}, {2, {28, 0, 28, 0}}, {4, {64, 0, 0, 0}}};
+    std::map<label, vec> exp_local_dim{
+        {1, {4, 4, 4, 4}}, {2, {8, 0, 8, 0}}, {4, {16, 0, 0, 0}}};
+
+    std::map<label, std::vector<label>> exp_non_local_nnz{
+        {1, {4, 4, 4, 4}}, {2, {4, 0, 4, 0}}, {4, {0, 0, 0, 0}}};
+
+    // expected local row indices in local indices
+    // first map fused true/false
+    // local row indices
+    /*
+     * The mesh has the following structure
+     *         local ids          |   global ids
+     * cells: [ 2 3 | 6 7 ]       |  [ 10 11 | 14 15 ]
+     *   2    [ 0 1 | 4 5 ]  3    |  [  8  9 | 12 13 ]
+     *        ------+------       |  --------+--------
+     *        [ 2 3 | 6 7 ]       |  [  2  3 |  6  7 ]
+     *   0    [ 0 1 | 4 5 ]  1    |  [  0  1 |  4  5 ]
+     *   */
+    vec nf_rows_2 = {0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4,
+                     4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 4, 6, 1, 3};
+    vec nf_cols_2 = {0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3, 4, 5,
+                     6, 4, 5, 7, 4, 6, 7, 5, 6, 7, 1, 3, 4, 6};
+    // [upper [0-3|12-15], lower [4,7|16-19], diag [8, 11|20-23], interfaces
+    // [24-27]]
+    vec nf_map_2 = {8,  0,  1,  4,  9,  2,
+                    5,  10, 3,  6,  7,  11,  // here first sd is done
+                    20, 12, 13, 16, 21, 14,
+                    17, 22, 15, 18, 19, 23,  // here interfaces start
+                    24, 25, 26, 27};
+    vec f_rows_2 = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3,
+                    4, 4, 4, 4, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7};
+    vec f_cols_2 = {0, 1, 2, 0, 1, 3, 4, 0, 2, 3, 1, 2, 3, 6,
+                    1, 4, 5, 6, 4, 5, 7, 3, 4, 6, 7, 5, 6, 7};
+
+    // [ rank 0            | rank 1             ] [ rank 0    | rank 1 ]
+    // [upper, lower, diag | upper, lower,  diag] [interfaces | interfaces]
+    // [0-3,   4-7,   8-11 | 12-15, 16-19, 20-23] [24-25,     | 26-27]
+    // interfaces start at 24
+    vec f_map_2 = {8,  0,  1,  4,  9,  2,  26, 5,  10, 3,  6,  7,  11, 27,
+                   24, 20, 12, 13, 16, 21, 14, 25, 17, 22, 15, 18, 19, 23};
+    /*
+     * The mesh has the following structure
+     *         local ids          |   global ids
+     * cells: [10 11|14 15]       |  [ 10 11 | 14 15 ]
+     *   2    [ 8 9 |12 13]  3    |  [  8  9 | 12 13 ]
+     *        ------+------       |  --------+--------
+     *        [ 2 3 | 6 7 ]       |  [  2  3 |  6  7 ]
+     *   0    [ 0 1 | 4 5 ]  1    |  [  0  1 |  4  5 ]
+     *   */
+    // NOTE the interfaces are in order of the comm_target indices
+    // ie [0, 0, 0 (1), (1), 2 (2), 2 (3), 2 (2), 2 (3)]
+    vec nf_rows_4 = {0,  0,  0,  1,  1,  1,  2,  2,  2,  3,  3,  3,  4,
+                     4,  4,  5,  5,  5,  6,  6,  6,  7,  7,  7,  8,  8,
+                     8,  9,  9,  9,  10, 10, 10, 11, 11, 11, 12, 12, 12,
+                     13, 13, 13, 14, 14, 14, 15, 15, 15, 4,  6,  8,  9,
+                     1,  3,  12, 13, 2,  3,  12, 14, 6,  7,  9,  11};
+    vec nf_cols_4 = {0,  1,  2,  0,  1,  3,  0,  2,  3,  1,  2,  3,  4,
+                     5,  6,  4,  5,  7,  4,  6,  7,  5,  6,  7,  8,  9,
+                     10, 8,  9,  11, 8,  10, 11, 9,  10, 11, 12, 13, 14,
+                     12, 13, 15, 12, 14, 15, 13, 14, 15, 1,  3,  2,  3,
+                     4,  6,  6,  7,  8,  9,  9,  11, 12, 13, 12, 14};
+    vec nf_map_4{8,  0,  1,  4,  9,  2,  5,  10, 3,  6,  7,  11,  // 1st sd done
+                 20, 12, 13, 16, 21, 14, 17, 22, 15, 18, 19, 23,  // 2nd sd done
+                 32, 24, 25, 28, 33, 26, 29, 34, 27, 30, 31, 35,  // 3rd sd done
+                 44, 36, 37, 40, 45, 38, 41, 46, 39, 42, 43, 47,  // 4th sd done
+                 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
+                 60, 61, 62, 63};
+    // fused
+    vec f_rows_4 = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  2,  3,  3,  3,  3,  3,   //
+        4,  4,  4,  4,  5,  5,  5,  6,  6,  6,  6,  6,  7,  7,  7,  7,   //
+        8,  8,  8,  8,  9,  9,  9,  9,  9,  10, 10, 10, 11, 11, 11, 11,  //
+        12, 12, 12, 12, 12, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15   //
+    };
+    vec f_cols_4 = {0,  1,  2,  0,  1,  3,  4,  0,  2,  3,  8,  1,  2,
+                    3,  6,  9,  1,  4,  5,  6,  4,  5,  7,  3,  4,  6,
+                    7,  12, 5,  6,  7,  13, 2,  8,  9,  10, 3,  8,  9,
+                    11, 12, 8,  10, 11, 9,  10, 11, 14, 6,  9,  12, 13,
+                    14, 7,  12, 13, 15, 11, 12, 14, 15, 13, 14, 15};
+    vec f_map_4{
+        8,  0,  1,  4,  9,  2,  52, 5,  10, 3,  56, 6,  7,  11, 53, 57,  //
+        48, 20, 12, 13, 16, 21, 14, 49, 17, 22, 15, 60, 18, 19, 23, 61,  //
+        50, 32, 24, 25, 51, 28, 33, 26, 62, 29, 34, 27, 30, 31, 35, 63,  //
+        54, 58, 44, 36, 37, 55, 40, 45, 38, 59, 41, 46, 39, 42, 43, 47};
+
+    std::map<bool, std::map<label, vec_vec>> exp_local_rows{
+        {false,
+         {{1, {rows, rows, rows, rows}},
+          {2, {nf_rows_2, {}, nf_rows_2, {}}},
+          {4, {nf_rows_4, {}, {}, {}}}}},
+        {true,
+         {{1, {rows, rows, rows, rows}},
+          {2, {f_rows_2, {}, f_rows_2, {}}},
+          {4, {f_rows_4, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_local_cols{
+        {false,
+         {{1, {cols, cols, cols, cols}},
+          {2, {nf_cols_2, {}, nf_cols_2, {}}},
+          {4, {nf_cols_4, {}, {}, {}}}}},
+        {true,
+         {{1, {cols, cols, cols, cols}},
+          {2, {f_cols_2, {}, f_cols_2, {}}},
+          {4, {f_cols_4, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_local_mapping{
+        {false,
+         {{1, {mapping, mapping, mapping, mapping}},
+          {2, {nf_map_2, {}, nf_map_2, {}}},
+          {4, {nf_map_4, {}, {}, {}}}}},
+        {true,
+         {{1, {mapping, mapping, mapping, mapping}},
+          {2, {f_map_2, {}, f_map_2, {}}},
+          {4, {f_map_4, {}, {}, {}}}}}};
+
+    // [start,end) indices of interface/submatrices
+    std::map<bool, std::map<label, vec_vec>> exp_local_spans_begin{
+        {false,
+         {{1, {{0}, {0}, {0}, {0}}},
+          {2, {{0, 24, 26}, {}, {0, 24, 26}, {}}},
+          {4, {{0, 48, 50, 52, 54, 56, 58, 60, 62}, {}, {}, {}}}}},
+        {true,
+         {{1, {{0}, {0}, {0}, {0}}},
+          {2, {{0}, {}, {0}, {}}},
+          {4, {{0}, {}, {}, {}}}}}};
+
+    // [start,end) indices of interface/submatrices
+    std::map<bool, std::map<label, vec_vec>> exp_local_spans_end{
+        {false,
+         {{1, {{13}, {13}, {13}, {13}}},
+          {2, {{24, 26, 28}, {}, {24, 26, 28}, {}}},
+          {4, {{48, 50, 52, 54, 56, 58, 60, 62, 64}, {}, {}, {}}}}},
+        {true,
+         {{1, {{13}, {13}, {13}, {13}}},
+          {2, {{28}, {}, {28}, {}}},
+          {4, {{64}, {}, {}, {}}}}}};
+
+    // non local data
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_rows{
+        {false,
+         {{1, non_local_rows},
+          {2, {{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, {{1, 2, 3, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 0, 1, 2}}},
+          {2, {{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    // non local cols are in global indices
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_cols{
+        {false,
+         {{1, {{4, 6, 8, 9}, {1, 3, 12, 13}, {2, 3, 12, 14}, {6, 7, 9, 11}}},
+          {2, {{8, 9, 12, 13}, {}, {2, 3, 6, 7}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, {{4, 8, 6, 9}, {1, 3, 12, 13}, {2, 3, 12, 14}, {6, 9, 7, 11}}},
+          {2, {{8, 9, 12, 13}, {}, {2, 3, 6, 7}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+
+    std::map<bool, std::map<label, vec_vec>> exp_non_local_mapping{
+        {false,
+         {{1, {{0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}}},
+          {2, {{0, 1, 2, 3}, {}, {0, 1, 2, 3}, {}}},
+          {4, {{}, {}, {}, {}}}}},
+        {true,
+         {{1, {{0, 2, 1, 3}, {0, 1, 2, 3}, {0, 1, 2, 3}, {0, 2, 1, 3}}},
+          {2, {{0, 1, 2, 3}, {}, {0, 1, 2, 3}, {}}},
+          {4, {{}, {}, {}, {}}}}}};
+};
+
+INSTANTIATE_TEST_SUITE_P(RepartitionerFixture2DInstantiation,
+                         RepartitionerFixture2D,
+                         testing::Combine(testing::Values(false, true),
+                                          testing::Values(1, 2, 4)),
+                         [](const auto &info) {
+                             // Can use info.param here to generate the test
+                             // suffix
+                             std::vector<std::string> names;
+                             names.emplace_back("fuse");
+                             names.emplace_back("ranks");
+                             std::string name = "fused_";
+                             name += std::to_string(std::get<0>(info.param));
+                             name += "_ranks_";
+                             name += std::to_string(std::get<1>(info.param));
+                             return name;
+                         });
+
+TEST_P(RepartitionerFixture2D, can_convert_to_global)
+{
+    // Arrange
+    using namespace gko::experimental::distributed;
+    auto ref_exec = exec.get_ref_exec();
+    auto partition = gko::share(
+        build_partition_from_local_size<label, label>(ref_exec, comm, 4));
+
+    vec_vec exp_global_idxs{
+        {4, 6, 8, 9},    // rank 0
+        {1, 3, 12, 13},  // rank 1
+        {2, 3, 12, 14},  // rank 2
+        {6, 7, 9, 11}    // rank 3
+    };
+
+    // Act
+    auto global_idxs =
+        detail::convert_to_global(partition, idxs[rank].data(),
+                                  non_local_spans[rank], comm_target_ids[rank]);
+
+    // Assert
+    EXPECT_EQ(global_idxs, exp_global_idxs[rank]);
+}
+
+TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
+{
+    // Arrange
+    auto [fused, ranks_per_gpu] = GetParam();
+    auto repartitioner =
+        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
+    auto ref_exec = exec.get_ref_exec();
+
+    // expected communication ranks
+    std::map<label, vec_vec> exp_res_ids{};
+    exp_res_ids[1] =
+        comm_target_ids;  // in the ranks_per_gpu==1 case nothing changes
+    // only communication partners are 0-2 and 2-0
+    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
+    // no communication if all ranks are repartitioned to single owner
+    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // expected communication sizes
+    std::map<label, vec_vec> exp_res_sizes;
+    exp_res_sizes.emplace(1, vec_vec({{2, 2}, {2, 2}, {2, 2}, {2, 2}}));
+    exp_res_sizes.emplace(2, vec_vec({{4}, {}, {4}, {}}));
+    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
+
+    // expected rows
+    vec_vec_vec rows{
+        {{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
+
+    std::map<label, vec_vec> exp_res_rows;
+    exp_res_rows.emplace(
+        1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}});
+    // after repartitioning with 2 ranks_per_gpu
+    // cells: [ 2 3  6 7 ]
+    //   2    [ 0 1  4 5 ]  3
+    //        -------------
+    //        [ 2 3  6 7 ]
+    //   0    [ 0 1  4 5 ]  1
+    exp_res_rows.emplace(2, vec_vec{{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}});
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <-
+    // local row ids
+    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    std::map<label, vec_vec> exp_gather_idx;
+    exp_gather_idx.emplace(
+        1, vec_vec{{0, 2, 0, 1}, {1, 3, 0, 1}, {2, 3, 0, 2}, {2, 3, 1, 3}});
+    exp_gather_idx.emplace(2, vec_vec{{0, 1, 4, 5}, {}, {2, 3, 6, 7}, {}});
+    exp_gather_idx.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // the original comm_pattern
+    auto comm_pattern = std::make_shared<CommunicationPattern>(
+        exec, comm_target_ids[rank], rows[rank]);
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
+
+    // Assert
+    auto res_ids = repart_comm_pattern->target_ids;
+
+    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
+
+    auto res_sizes = repart_comm_pattern->target_sizes;
+    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+
+    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    auto res_rows = total_rank_send_idx;
+    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
+
+    auto total_recv_gather_idx =
+        repart_comm_pattern->compute_recv_gather_idxs(exec);
+    auto res_gather_idx = convert_to_vector(total_recv_gather_idx);
+    EXPECT_EQ(res_gather_idx, exp_gather_idx[ranks_per_gpu][rank]);
+}
+
+TEST_P(RepartitionerFixture2D, can_repartition_sparsity_pattern)
 {
     // Arrange
-    label ranks_per_gpu = GetParam();
-    bool fused = false;
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto [fused, ranks_per_gpu] = GetParam();
+    auto repartitioner =
+        Repartitioner(local_size, ranks_per_gpu, 0, exec, fused);
     auto ref_exec = exec.get_ref_exec();
 
     std::vector<label> ranks{rank};
     auto local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{4, 4}, rows, cols, mapping, spans, ranks);
+        ref_exec, gko::dim<2>{4, 4}, rows, cols, mapping, spans);
 
     auto non_local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{4, non_local_ranks[rank].size()},
-        non_local_rows[rank], non_local_cols[rank], non_local_mapping[rank],
-        non_local_spans[rank], non_local_ranks[rank]);
+        ref_exec, gko::dim<2>{4, 4}, non_local_rows[rank], non_local_cols[rank],
+        non_local_mapping[rank], non_local_spans[rank]);
 
     // Act
     auto [repart_local, repart_non_local, tracking] =
         repartitioner.repartition_sparsity(exec, local_sparsity,
-                                           non_local_sparsity, fused);
+                                           non_local_sparsity,
+                                           comm_target_ids[rank], fused);
     // Assert
     // local properties
     ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[0], exp_local_dim_rows[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_local->dim[1], exp_local_dim_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[1], exp_local_dim[ranks_per_gpu][rank]);
 
     auto res_local_rows = convert_to_vector(repart_local->row_idxs);
     auto res_local_cols = convert_to_vector(repart_local->col_idxs);
     auto res_local_mapping = convert_to_vector(repart_local->ldu_mapping);
-    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_cols, exp_local_cols[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_mapping, exp_local_mapping[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_rows.size(),
+              exp_local_rows[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_rows.size(); i++) {
+        ASSERT_EQ(res_local_rows[i],
+                  exp_local_rows[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i;
+    }
+    ASSERT_EQ(res_local_cols.size(),
+              exp_local_cols[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_cols.size(); i++) {
+        ASSERT_EQ(res_local_cols[i],
+                  exp_local_cols[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i;
+    }
+    ASSERT_EQ(res_local_mapping.size(),
+              exp_local_mapping[fused][ranks_per_gpu][rank].size());
+    for (size_t i = 0; i < res_local_mapping.size(); i++) {
+        ASSERT_EQ(res_local_mapping[i],
+                  exp_local_mapping[fused][ranks_per_gpu][rank][i])
+            << " failed at index " << i;
+    }
 
     std::vector<label> res_local_spans_begin{};
     std::vector<label> res_local_spans_end{};
@@ -43,18 +404,24 @@ TEST_P(RepartitionerFixture2D, can_repartition_sparsity_pattern_2D_for_n_ranks_u
         res_local_spans_end.push_back(end);
     }
     ASSERT_EQ(res_local_spans_begin,
-              exp_local_spans_begin[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_spans_end, exp_local_spans_end[ranks_per_gpu][rank]);
+              exp_local_spans_begin[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_spans_end,
+              exp_local_spans_end[fused][ranks_per_gpu][rank]);
 
     // non local properties
     ASSERT_EQ(repart_non_local->num_nnz,
               exp_non_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[0],
-              exp_local_dim_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim[ranks_per_gpu][rank]);
     ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
 
     auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
     auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);
-    ASSERT_EQ(res_non_local_rows, exp_non_local_rows[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_non_local_cols, exp_non_local_cols[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_rows,
+              exp_non_local_rows[fused][ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_cols,
+              exp_non_local_cols[fused][ranks_per_gpu][rank]);
+    auto res_non_local_mapping =
+        convert_to_vector(repart_non_local->ldu_mapping);
+    ASSERT_EQ(res_non_local_mapping,
+              exp_non_local_mapping[fused][ranks_per_gpu][rank]);
 }