[NDTensorsMappedArraysExt] Support for using MappedArrays as data of …

…ITensors (#1505)

* [NDTensorsMappedArraysExt] Support for using MappedArrays as data of ITensors

* [NDTensors] Bump to v0.3.32

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.31"
+version = "0.3.32"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -36,6 +36,7 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
+MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
@@ -46,6 +47,7 @@ NDTensorsAMDGPUExt = ["AMDGPU", "GPUArraysCore"]
 NDTensorsCUDAExt = ["CUDA", "GPUArraysCore"]
 NDTensorsGPUArraysCoreExt = "GPUArraysCore"
 NDTensorsHDF5Ext = "HDF5"
+NDTensorsMappedArraysExt = ["MappedArrays"]
 NDTensorsMetalExt = ["GPUArraysCore", "Metal"]
 NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"

NDTensors/ext/NDTensorsMappedArraysExt/NDTensorsMappedArraysExt.jl

@@ -0,0 +1,25 @@
+module NDTensorsMappedArraysExt
+using MappedArrays: AbstractMappedArray
+using NDTensors: NDTensors
+function NDTensors.similar(arraytype::Type{<:AbstractMappedArray}, dims::Tuple{Vararg{Int}})
+  return similar(Array{eltype(arraytype)}, dims)
+end
+function NDTensors.similartype(storagetype::Type{<:AbstractMappedArray})
+  return Array{eltype(storagetype),ndims(storagetype)}
+end
+function NDTensors.similartype(
+  storagetype::Type{<:AbstractMappedArray}, dims::Tuple{Vararg{Int}}
+)
+  return Array{eltype(storagetype),length(dims)}
+end
+
+using MappedArrays: ReadonlyMappedArray
+using NDTensors: AllowAlias
+# It is a bit unfortunate that we have to define this, it fixes an ambiguity
+# error with MappedArrays.
+function (arraytype::Type{ReadonlyMappedArray{T,N,A,F}} where {T,N,A<:AbstractArray,F})(
+  ::AllowAlias, a::AbstractArray
+)
+  return a
+end
+end

NDTensors/src/lib/TypeParameterAccessors/src/base/abstractarray.jl

@@ -68,6 +68,14 @@ end
   return set_type_parameter(type, eltype, param)
 end
 
+# These are generic fallback definitions. By convention,
+# this is very commonly true of `AbstractArray` subtypes
+# but it may not be correct, but it is very convenient
+# to define this to make more operations "just work"
+# on most AbstractArrays.
+position(type::Type{<:AbstractArray}, ::typeof(eltype)) = Position(1)
+position(type::Type{<:AbstractArray}, ::typeof(ndims)) = Position(2)
+
 for wrapper in [:PermutedDimsArray, :(Base.ReshapedArray), :SubArray]
   @eval begin
     position(type::Type{<:$wrapper}, ::typeof(eltype)) = Position(1)

NDTensors/src/lib/TypeParameterAccessors/test/test_basics.jl

@@ -15,6 +15,11 @@ using NDTensors.TypeParameterAccessors:
 include("utils/test_inferred.jl")
 @testset "TypeParameterAccessors basics" begin
   @testset "Get parameters" begin
+    @test_inferred type_parameter(AbstractArray{Float64}, 1) == Float64 wrapped = true
+    @test_inferred type_parameter(AbstractArray{Float64}, Position(1)) == Float64
+    @test_inferred type_parameter(AbstractArray{Float64}, eltype) == Float64
+    @test_inferred type_parameter(AbstractMatrix{Float64}, ndims) == 2
+
     @test_inferred type_parameter(Array{Float64}, 1) == Float64 wrapped = true
     @test_inferred type_parameter(Array{Float64}, Position(1)) == Float64
     @test_inferred type_parameter(Val{3}) == 3

NDTensors/test/Project.toml

@@ -1,15 +1,16 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -22,11 +23,11 @@ TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
-[extras]
-AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
-Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
-
 [compat]
 Metal = "1.1.0"
 cuTENSOR = "2.0"
+
+[extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"

NDTensors/test/ext/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
+MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"

NDTensors/test/ext/runtests.jl

@@ -0,0 +1,12 @@
+@eval module $(gensym())
+using Test: @testset
+@testset "$(@__DIR__)" begin
+  filenames = filter(readdir(@__DIR__)) do f
+    startswith("test_")(f) && endswith(".jl")(f)
+  end
+  @testset "Test $(@__DIR__)/$filename" for filename in filenames
+    println("Running $(@__DIR__)/$filename")
+    include(filename)
+  end
+end
+end

NDTensors/test/ext/test_mappedarraysext.jl

@@ -0,0 +1,24 @@
+@eval module $(gensym())
+using ITensors: Index, itensor
+using LinearAlgebra: qr, svd
+using MappedArrays: mappedarray
+using Test: @test, @testset
+f(i::Int...) = float(sum(iⱼ -> iⱼ^2, i))
+f(i::CartesianIndex) = f(Tuple(i)...)
+@testset "NDTensorsMappedArraysExt" begin
+  a = mappedarray(f, CartesianIndices((2, 2)))
+  b = copy(a)
+  i, j = Index.((2, 2))
+  ta = itensor(a, i, j)
+  tb = itensor(b, i, j)
+  @test ta ≈ tb
+  @test ta[i => 1, j => 2] ≈ tb[i => 1, j => 2]
+  @test 2 * ta ≈ 2 * tb
+  @test ta + ta ≈ tb + tb
+  @test ta * ta ≈ tb * tb
+  ua, sa, va = svd(ta, i)
+  @test ua * sa * va ≈ ta
+  qa, ra = qr(ta, i)
+  @test qa * ra ≈ ta
+end
+end

NDTensors/test/runtests.jl

@@ -7,7 +7,7 @@ using SafeTestsets: @safetestset
     filenames = filter(readdir(@__DIR__)) do f
       startswith("test_")(f) && endswith(".jl")(f)
     end
-    for dir in ["lib", "ITensors"]
+    for dir in ["lib", "ext", "ITensors"]
       push!(filenames, joinpath(dir, "runtests.jl"))
     end
     @testset "Test $(@__DIR__)/$filename" for filename in filenames