diff --git a/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysAdaptExt/src/BlockSparseArraysAdaptExt.jl b/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysAdaptExt/src/BlockSparseArraysAdaptExt.jl new file mode 100644 index 0000000000..68cbf05e35 --- /dev/null +++ b/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 diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl index 3542c3e10b..af59ae7f51 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl +++ b/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 diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/map.jl new file mode 100644 index 0000000000..7ecad61a41 --- /dev/null +++ b/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