Sparse array slicing

ITensor · Nov 13, 2024 · c98436b · c98436b
1 parent cf050da
commit c98436b
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl
@@ -143,6 +143,17 @@ function sparse_setindex!(a::AbstractArray, value, I::CartesianIndex{1})
   return a
 end
 
+# Slicing
+function sparse_setindex!(a::AbstractArray, value, I::AbstractUnitRange...)
+  inds = CartesianIndices(I)
+  for i in stored_indices(value)
+    if i in CartesianIndices(inds)
+      a[inds[i]] = value[i]
+    end
+  end
+  return a
+end
+
 # Handle trailing indices
 function sparse_setindex!(a::AbstractArray, value, I::CartesianIndex)
   t = Tuple(I)