Adapt export

FEniCS · Aug 19, 2024 · e9cf2c9 · e9cf2c9
1 parent 1c6c8e0
commit e9cf2c9
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/python/dolfinx/wrappers/la.cpp b/python/dolfinx/wrappers/la.cpp
@@ -135,10 +135,10 @@ void declare_objects(nb::module_& m, const std::string& type)
            {
              const std::array<int, 2> bs = self.block_size();
              std::size_t nrows = self.num_all_rows() * bs[0];
-             auto map_col = self.index_map(1);
-             std::size_t ncols
-                 = (map_col->size_local() + map_col->num_ghosts()) * bs[1];
-             return dolfinx_wrappers::as_nbarray(self.to_dense(),
+             std::size_t ncols = self.index_map(1)->size_global() * bs[1];
+             auto dense = self.to_dense();
+             assert(nrows*ncols == dense.size());
+             return dolfinx_wrappers::as_nbarray(std::move(self.to_dense()),
                                                  {nrows, ncols});
            })
       .def_prop_ro(