Fix some issues introduces in map

ITensor · Jun 3, 2024 · aca628f · aca628f
1 parent 5e987b8
commit aca628f
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl
@@ -32,8 +32,14 @@ function SparseArrayInterface.sparse_map!(
   for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...)
     BI_dest = blockindexrange(a_dest, I)
     BI_srcs = map(a_src -> blockindexrange(a_src, I), a_srcs)
-    block_dest = @view a_dest[_block(BI_dest)]
-    block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs))
+    # TODO: Investigate why this doesn't work:
+    # block_dest = @view a_dest[_block(BI_dest)]
+    block_dest = blocks(a_dest)[Int.(Tuple(_block(BI_dest)))...]
+    # TODO: Investigate why this doesn't work:
+    # block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs))
+    block_srcs = ntuple(length(a_srcs)) do i
+      return blocks(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...]
+    end
     subblock_dest = @view block_dest[BI_dest.indices...]
     subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs))
     # TODO: Use `map!!` to handle immutable blocks.