Add adapt support to BlockSparseArrays

Author: mtfishman

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysAdaptExt/src/BlockSparseArraysAdaptExt.jl

@@ -0,0 +1,5 @@
+module BlockSparseArraysAdaptExt
+using Adapt: Adapt, adapt
+using ..BlockSparseArrays: AbstractBlockSparseArray, map_stored_blocks
+Adapt.adapt_structure(to, x::AbstractBlockSparseArray) = map_stored_blocks(adapt(to), x)
+end

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

@@ -4,6 +4,7 @@ include("blocksparsearrayinterface/blocksparsearrayinterface.jl")
 include("blocksparsearrayinterface/linearalgebra.jl")
 include("blocksparsearrayinterface/blockzero.jl")
 include("blocksparsearrayinterface/broadcast.jl")
+include("blocksparsearrayinterface/map.jl")
 include("blocksparsearrayinterface/arraylayouts.jl")
 include("blocksparsearrayinterface/views.jl")
 include("abstractblocksparsearray/abstractblocksparsearray.jl")
@@ -20,4 +21,5 @@ include("blocksparsearray/blocksparsearray.jl")
 include("BlockArraysSparseArrayInterfaceExt/BlockArraysSparseArrayInterfaceExt.jl")
 include("../ext/BlockSparseArraysTensorAlgebraExt/src/BlockSparseArraysTensorAlgebraExt.jl")
 include("../ext/BlockSparseArraysGradedAxesExt/src/BlockSparseArraysGradedAxesExt.jl")
+include("../ext/BlockSparseArraysAdaptExt/src/BlockSparseArraysAdaptExt.jl")
 end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/map.jl

@@ -0,0 +1,14 @@
+function map_stored_blocks(f, a::AbstractArray)
+  # TODO: Implement this as:
+  # ```julia
+  # mapped_blocks = SparseArraysInterface.map_stored(f, blocks(a))
+  # BlockSparseArray(mapped_blocks, axes(a))
+  # ```
+  # TODO: `block_stored_indices` should output `Indices` storing
+  # the stored Blocks, not a `Dictionary` from cartesian indices
+  # to Blocks.
+  bs = block_stored_indices(a)
+  mapped_blocks = Dictionary(bs, map(b -> f(@view(a[b])), bs))
+  # TODO: Use `similartype(typeof(a), eltype(eltype(mapped_blocks)))(...)`.
+  return BlockSparseArray(mapped_blocks, axes(a))
+end