Fix slicing on GPU

Author: mtfishman

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -1,3 +1,4 @@
+using ArrayLayouts: ArrayLayouts, MemoryLayout, sub_materialize
 using BlockArrays:
   BlockArrays,
   AbstractBlockArray,
@@ -537,6 +538,20 @@ function SparseArrayInterface.nstored(a::BlockView)
   return 0
 end
 
+## # Allow more fine-grained control:
+## function ArrayLayouts.sub_materialize(layout, a::BlockView, ax)
+##   return blocks(a.array)[Int.(a.block)...]
+## end
+## function ArrayLayouts.sub_materialize(layout, a::BlockView)
+##   return sub_materialize(layout, a, axes(a))
+## end
+## function ArrayLayouts.sub_materialize(a::BlockView)
+##   return sub_materialize(MemoryLayout(a), a)
+## end
+function ArrayLayouts.sub_materialize(a::BlockView)
+  return blocks(a.array)[Int.(a.block)...]
+end
+
 function view!(a::AbstractArray{<:Any,N}, index::Block{N}) where {N}
   return view!(a, Tuple(index)...)
 end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/arraylayouts.jl

@@ -22,8 +22,10 @@ end
 
 # Materialize a SubArray view.
 function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, axes)
-  # TODO: Make more generic for GPU.
-  a_dest = BlockSparseArray{eltype(a)}(axes)
+  # TODO: Define `blocktype`/`blockstype` for `SubArray` wrapping `BlockSparseArray`.
+  # TODO: Use `similar`?
+  blocktype_a = blocktype(parent(a))
+  a_dest = BlockSparseArray{eltype(a),length(axes),blocktype_a}(axes)
   a_dest .= a
   return a_dest
 end
@@ -32,8 +34,7 @@ end
 function ArrayLayouts.sub_materialize(
   layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}}
 )
-  # TODO: Make more generic for GPU.
-  a_dest = Array{eltype(a)}(undef, length.(axes))
+  a_dest = blocktype(a)(undef, length.(axes))
   a_dest .= a
   return a_dest
 end