Define a generic dense function (#72)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseArraysBase"
 uuid = "0d5efcca-f356-4864-8770-e1ed8d78f208"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.7.3"
+version = "0.7.4"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -10,9 +10,11 @@ ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [compat]
 Accessors = "0.1.41"
@@ -22,12 +24,14 @@ ArrayLayouts = "1.11.0"
 DerivableInterfaces = "0.5"
 Dictionaries = "0.4.3"
 FillArrays = "1.13.0"
+GPUArraysCore = "0.2.0"
 LinearAlgebra = "1.10"
 MapBroadcast = "0.1.5"
 Random = "1.10.0"
 SafeTestsets = "0.1"
 Suppressor = "0.2"
 Test = "1.10"
+TypeParameterAccessors = "0.4.3"
 julia = "1.10"
 
 [extras]

src/abstractsparsearrayinterface.jl

@@ -39,15 +39,17 @@ unstoredsimilar(a::AbstractArray) = a
 # Generic functionality for converting to a
 # dense array, trying to preserve information
 # about the array (such as which device it is on).
-# TODO: Maybe call `densecopy`?
-# TODO: Make sure this actually preserves the device,
-# maybe use `TypeParameterAccessors.unwrap_array_type`.
-# TODO: Turn into an `@interface` function.
-function densearray(a::AbstractArray)
-  # TODO: `set_ndims(unwrap_array_type(a), ndims(a))(a)`
-  # Maybe define `densetype(a) = set_ndims(unwrap_array_type(a), ndims(a))`.
-  # Or could use `unspecify_parameters(unwrap_array_type(a))(a)`.
-  return Array(a)
+using TypeParameterAccessors: unspecify_type_parameters, unwrap_array, unwrap_array_type
+function densetype(arraytype::Type{<:AbstractArray})
+  return unspecify_type_parameters(unwrap_array_type(arraytype))
+end
+# TODO: Ideally this would be defined as `densetype(typeof(a))` but that is less general right now since `unwrap_array_type` is defined on fewer arrays, since it is based on `parentype` rather than `parent`.
+function densetype(a::AbstractArray)
+  return unspecify_type_parameters(typeof(unwrap_array(a)))
+end
+using GPUArraysCore: @allowscalar
+function dense(a::AbstractArray)
+  return @allowscalar convert(densetype(a), a)
 end
 
 # Minimal interface for `SparseArrayInterface`.

src/sparsearraydok.jl

@@ -122,3 +122,8 @@ end
 function ArrayLayouts.zero!(a::SparseArrayDOK)
   return zero!(a)
 end
+
+# The fallback doesn't work since `SparseArrayDOK`
+# isn't a wrapper around `Array`.
+densetype(::SparseArrayDOK) = Array
+densetype(::Type{<:SparseArrayDOK}) = Array

test/test_dense.jl

@@ -0,0 +1,41 @@
+using Adapt: adapt
+using JLArrays: JLArray
+using SparseArraysBase: dense, sparsezeros
+using Test: @test, @testset
+
+elts = (Float32, ComplexF64)
+arrayts = (Array, JLArray)
+@testset "dense (arraytype=$arrayt, eltype=$elt)" for arrayt in arrayts, elt in elts
+  dev(x) = adapt(arrayt, x)
+
+  @testset "SparseArrayDOK" begin
+    s = sparsezeros(elt, 3, 4)
+    s[1, 2] = 2
+    d = dense(s)
+    @test d isa Matrix{elt}
+    @test d == [0 2 0 0; 0 0 0 0; 0 0 0 0]
+  end
+
+  @testset "Custom sparse array" begin
+    struct MySparseArrayDOK{T,N,S<:AbstractVector{T}} <: AbstractArray{T,N}
+      storedvalues::S
+      storedindices::Dict{CartesianIndex{N},Int}
+      size::NTuple{N,Int}
+    end
+    Base.size(a::MySparseArrayDOK) = a.size
+    function Base.getindex(a::MySparseArrayDOK{<:Any,N}, I::Vararg{Int,N}) where {N}
+      storageindex = get(a.storedindices, CartesianIndex(I), nothing)
+      isnothing(storageindex) && return zero(eltype(a))
+      return a.storedvalues[storageindex]
+    end
+    Base.parent(a::MySparseArrayDOK) = a.storedvalues
+
+    s = MySparseArrayDOK(
+      dev(elt[2, 4]), Dict([CartesianIndex(1, 2) => 1, CartesianIndex(3, 4) => 2]), (3, 4)
+    )
+    d = dense(s)
+    @show typeof(d)
+    @test d isa arrayt{elt,2}
+    @test d == dev(elt[0 2 0 0; 0 0 0 0; 0 0 0 4])
+  end
+end