Merge branch 'main' into non-abelian_fusion

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.13"
+version = "0.3.18"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -14,7 +14,6 @@ EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -35,17 +34,19 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
-NDTensorsAMDGPUExt = "AMDGPU"
-NDTensorsCUDAExt = "CUDA"
+NDTensorsAMDGPUExt = ["AMDGPU","GPUArraysCore"]
+NDTensorsCUDAExt = ["CUDA","GPUArraysCore"]
+NDTensorsGPUArraysCoreExt = "GPUArraysCore"
 NDTensorsHDF5Ext = "HDF5"
-NDTensorsMetalExt = "Metal"
+NDTensorsMetalExt = ["GPUArraysCore", "Metal"]
 NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"
 NDTensorscuTENSORExt = "cuTENSOR"
@@ -90,6 +91,7 @@ julia = "1.6"
 [extras]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"

NDTensors/ext/NDTensorsAMDGPUExt/append.jl

@@ -1,5 +1,5 @@
-using GPUArraysCore: @allowscalar
 using AMDGPU: ROCArray
+using GPUArraysCore: @allowscalar
 using NDTensors.Expose: Exposed, unexpose
 
 ## Warning this append function uses scalar indexing and is therefore extremely slow

NDTensors/ext/NDTensorsAMDGPUExt/indexing.jl

@@ -1,7 +1,7 @@
-using NDTensors.Expose: Exposed, expose, parent, unexpose
-using NDTensors.GPUArraysCoreExtensions: cpu
 using AMDGPU: AMDGPU, ROCArray
 using GPUArraysCore: @allowscalar
+using NDTensors.Expose: Exposed, expose, parent, unexpose
+using NDTensors.GPUArraysCoreExtensions: cpu
 
 function Base.getindex(E::Exposed{<:ROCArray})
   return @allowscalar unexpose(E)[]

NDTensors/ext/NDTensorsCUDAExt/append.jl

@@ -1,5 +1,5 @@
-using GPUArraysCore: @allowscalar
 using CUDA: CuArray
+using GPUArraysCore: @allowscalar
 using NDTensors.Expose: Exposed, unexpose
 
 ## Warning this append function uses scalar indexing and is therefore extremely slow

NDTensors/ext/NDTensorsGPUArraysCoreExt/NDTensorsGPUArraysCoreExt.jl

@@ -0,0 +1,3 @@
+module NDTensorsGPUArraysCoreExt
+include("contract.jl")
+end

NDTensors/ext/NDTensorsGPUArraysCoreExt/contract.jl

@@ -0,0 +1,48 @@
+using Adapt: adapt
+using GPUArraysCore: AbstractGPUArray
+using NDTensors: NDTensors, DenseTensor, DiagTensor, contract!, dense, inds, Tensor
+using NDTensors.Expose: Exposed, expose, unexpose
+using NDTensors.TypeParameterAccessors: parenttype, set_ndims
+
+## In this function we convert the DiagTensor to a dense tensor and
+## Feed it back into contract
+function NDTensors.contract!(
+  output_tensor::Exposed{<:AbstractGPUArray,<:DenseTensor},
+  labelsoutput_tensor,
+  tensor1::Exposed{<:Number,<:DiagTensor},
+  labelstensor1,
+  tensor2::Exposed{<:AbstractGPUArray,<:DenseTensor},
+  labelstensor2,
+  α::Number=one(Bool),
+  β::Number=zero(Bool),
+)
+  tensor1 = unexpose(tensor1)
+  ## convert tensor1 to a dense
+  ## TODO this allocates on CPU first then moves over to GPU which could be slow
+  tensor1 = adapt(set_ndims(parenttype(typeof(tensor2)), 1), dense(tensor1))
+  return contract!(
+    output_tensor,
+    labelsoutput_tensor,
+    expose(tensor1),
+    labelstensor1,
+    tensor2,
+    labelstensor2,
+    α,
+    β,
+  )
+end
+
+function NDTensors.contract!(
+  output_tensor::Exposed{<:AbstractGPUArray,<:DenseTensor},
+  labelsoutput_tensor,
+  tensor1::Exposed{<:AbstractGPUArray,<:DenseTensor},
+  labelstensor1,
+  tensor2::Exposed{<:Number,<:DiagTensor},
+  labelstensor2,
+  α::Number=one(Bool),
+  β::Number=zero(Bool),
+)
+  return contract!(
+    output_tensor, labelsoutput_tensor, tensor2, labelstensor2, tensor1, labelstensor1, α, β
+  )
+end

NDTensors/ext/NDTensorsMetalExt/permutedims.jl

@@ -1,4 +1,5 @@
 using Metal: MtlArray
+using GPUArraysCore: @allowscalar
 using NDTensors.Expose: Exposed, expose, unexpose
 ## Theres an issue in metal that `ReshapedArray' wrapped arrays cannot be permuted using
 ## permutedims (failing in that Metal uses scalar indexing)

NDTensors/src/blocksparse/diagblocksparse.jl

@@ -639,6 +639,11 @@ function contract!(
   labelsT2,
   contraction_plan,
 ) where {ElR<:Number,NR}
+  if any(b -> !allequal(Tuple(b)), nzblocks(T2))
+    return error(
+      "When contracting a BlockSparse tensor with a DiagBlockSparse tensor, the DiagBlockSparse tensor must be block diagonal for the time being.",
+    )
+  end
   already_written_to = Dict{Block{NR},Bool}()
   indsR = inds(R)
   indsT1 = inds(T1)
@@ -661,7 +666,9 @@ function contract!(
       # Overwrite the block of R
       β = zero(ElR)
     end
-    contract!(Rblock, labelsR, T1block, labelsT1, T2block, labelsT2, α, β)
+    contract!(
+      expose(Rblock), labelsR, expose(T1block), labelsT1, expose(T2block), labelsT2, α, β
+    )
   end
   return R
 end

NDTensors/src/imports.jl

@@ -9,7 +9,6 @@ using Base.Threads
 using Compat
 using Dictionaries
 using Folds
-using GPUArraysCore
 using InlineStrings
 using Random
 using LinearAlgebra

NDTensors/src/lib/BackendSelection/src/backend_types.jl

@@ -26,7 +26,7 @@ for type in (:Algorithm, :Backend)
     parameters(backend::$type) = getfield(backend, :kwargs)
 
     function Base.show(io::IO, backend::$type)
-      return print(io, "$type type ", backend_string(backend), ", ", parameters(backend))
+      return print(io, "$($type) type ", backend_string(backend), ", ", parameters(backend))
     end
     Base.print(io::IO, backend::$type) =
       print(io, backend_string(backend), ", ", parameters(backend))

NDTensors/src/lib/BackendSelection/test/runtests.jl

@@ -24,5 +24,6 @@ using NDTensors.AlgorithmSelection: AlgorithmSelection
   # Macro syntax.
   @test Algorithm"backend"(; x=2, y=3) === Algorithm("backend"; x=2, y=3)
   @test Backend"backend"(; x=2, y=3) === Backend("backend"; x=2, y=3)
+  @test isnothing(show(Algorithm("")))
 end
 end

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

@@ -1,8 +1,22 @@
 module BlockSparseArraysGradedAxesExt
-using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange
-using ..BlockSparseArrays: BlockSparseArrays, block_merge
+using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange, blocks
+using ..BlockSparseArrays:
+  BlockSparseArrays,
+  AbstractBlockSparseArray,
+  AbstractBlockSparseMatrix,
+  BlockSparseArray,
+  BlockSparseMatrix,
+  block_merge
 using ...GradedAxes:
-  GradedUnitRange, OneToOne, blockmergesortperm, blocksortperm, invblockperm, tensor_product
+  GradedUnitRange,
+  OneToOne,
+  blockmergesortperm,
+  blocksortperm,
+  dual,
+  invblockperm,
+  nondual,
+  tensor_product
+using LinearAlgebra: Adjoint, Transpose
 using ...TensorAlgebra:
   TensorAlgebra, FusionStyle, BlockReshapeFusion, SectorFusion, fusedims, splitdims
 
@@ -45,4 +59,68 @@ function TensorAlgebra.splitdims(
   a_blockpermed = a[blockperms...]
   return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...)
 end
+
+# This is a temporary fix for `eachindex` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.eachindex(a::AbstractBlockSparseArray)
+  return CartesianIndices(nondual.(axes(a)))
+end
+
+# TODO: Handle this through some kind of trait dispatch, maybe
+# a `SymmetryStyle`-like trait to check if the block sparse
+# matrix has graded axes.
+function Base.axes(a::Adjoint{<:Any,<:AbstractBlockSparseMatrix})
+  return dual.(reverse(axes(a')))
+end
+
+# This is a temporary fix for `show` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function blocksparse_show(
+  io::IO, mime::MIME"text/plain", a::AbstractArray, axes_a::Tuple; kwargs...
+)
+  println(io, "typeof(axes) = ", typeof(axes_a), "\n")
+  println(
+    io,
+    "Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.\n",
+  )
+  return invoke(show, Tuple{IO,MIME"text/plain",AbstractArray}, io, mime, a; kwargs...)
+end
+
+# This is a temporary fix for `show` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.show(io::IO, mime::MIME"text/plain", a::BlockSparseArray; kwargs...)
+  axes_a = axes(a)
+  a_nondual = BlockSparseArray(blocks(a), nondual.(axes(a)))
+  return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...)
+end
+
+# This is a temporary fix for `show` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.show(
+  io::IO, mime::MIME"text/plain", a::Adjoint{<:Any,<:BlockSparseMatrix}; kwargs...
+)
+  axes_a = axes(a)
+  a_nondual = BlockSparseArray(blocks(a'), dual.(nondual.(axes(a))))'
+  return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...)
+end
+
+# This is a temporary fix for `show` being broken for BlockSparseArrays
+# with mixed dual and non-dual axes. This shouldn't be needed once
+# GradedAxes is rewritten using BlockArrays v1.
+# TODO: Delete this once GradedAxes is rewritten.
+function Base.show(
+  io::IO, mime::MIME"text/plain", a::Transpose{<:Any,<:BlockSparseMatrix}; kwargs...
+)
+  axes_a = axes(a)
+  a_nondual = tranpose(BlockSparseArray(transpose(blocks(a)), nondual.(axes(a))))
+  return blocksparse_show(io, mime, a_nondual, axes_a; kwargs...)
+end
 end

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl

@@ -3,9 +3,8 @@ using Compat: Returns
 using Test: @test, @testset, @test_broken
 using BlockArrays: Block, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
-using NDTensors.GradedAxes: GradedUnitRange, gradedrange
+using NDTensors.GradedAxes: GradedAxes, GradedUnitRange, UnitRangeDual, dual, gradedrange
 using NDTensors.LabelledNumbers: label
-using NDTensors.Sectors: U1
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
 using Random: randn!
@@ -16,6 +15,14 @@ function blockdiagonal!(f, a::AbstractArray)
   end
   return a
 end
+
+struct U1
+  n::Int
+end
+GradedAxes.dual(c::U1) = U1(-c.n)
+GradedAxes.fuse_labels(c1::U1, c2::U1) = U1(c1.n + c2.n)
+Base.isless(c1::U1, c2::U1) = isless(c1.n, c2.n)
+
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "BlockSparseArraysGradedAxesExt (eltype=$elt)" for elt in elts
   @testset "map" begin
@@ -70,5 +77,38 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test_broken blocksize(m) == (3, 3)
     @test a == splitdims(m, (d1, d2), (d1, d2))
   end
+  @testset "dual axes" begin
+    r = gradedrange([U1(0) => 2, U1(1) => 2])
+    a = BlockSparseArray{elt}(dual(r), r)
+    a[Block(1, 1)] = randn(size(a[Block(1, 1)]))
+    a[Block(2, 2)] = randn(size(a[Block(2, 2)]))
+    a_dense = Array(a)
+    @test eachindex(a) == CartesianIndices(size(a))
+    for I in eachindex(a)
+      @test a[I] == a_dense[I]
+    end
+    @test axes(a') == dual.(reverse(axes(a)))
+    # TODO: Define and use `isdual` here.
+    @test axes(a', 1) isa UnitRangeDual
+    @test !(axes(a', 2) isa UnitRangeDual)
+    @test isnothing(show(devnull, MIME("text/plain"), a))
+  end
+  @testset "Matrix multiplication" begin
+    r = gradedrange([U1(0) => 2, U1(1) => 3])
+    a1 = BlockSparseArray{elt}(dual(r), r)
+    a1[Block(1, 2)] = randn(elt, size(@view(a1[Block(1, 2)])))
+    a1[Block(2, 1)] = randn(elt, size(@view(a1[Block(2, 1)])))
+    a2 = BlockSparseArray{elt}(dual(r), r)
+    a2[Block(1, 2)] = randn(elt, size(@view(a2[Block(1, 2)])))
+    a2[Block(2, 1)] = randn(elt, size(@view(a2[Block(2, 1)])))
+    @test Array(a1 * a2) ≈ Array(a1) * Array(a2)
+    @test Array(a1' * a2') ≈ Array(a1') * Array(a2')
+
+    a2 = BlockSparseArray{elt}(r, dual(r))
+    a2[Block(1, 2)] = randn(elt, size(@view(a2[Block(1, 2)])))
+    a2[Block(2, 1)] = randn(elt, size(@view(a2[Block(2, 1)])))
+    @test Array(a1' * a2) ≈ Array(a1') * Array(a2)
+    @test Array(a1 * a2') ≈ Array(a1) * Array(a2')
+  end
 end
 end

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

@@ -11,7 +11,6 @@ include("abstractblocksparsearray/abstractblocksparsevector.jl")
 include("abstractblocksparsearray/view.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/sparsearrayinterface.jl")
-include("abstractblocksparsearray/linearalgebra.jl")
 include("abstractblocksparsearray/broadcast.jl")
 include("abstractblocksparsearray/map.jl")
 include("blocksparsearray/defaults.jl")

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

@@ -1,7 +1,6 @@
 using ArrayLayouts: ArrayLayouts, MemoryLayout, MulAdd
 using BlockArrays: BlockLayout
 using ..SparseArrayInterface: SparseLayout
-using LinearAlgebra: mul!
 
 function ArrayLayouts.MemoryLayout(arraytype::Type{<:BlockSparseArrayLike})
   outer_layout = typeof(MemoryLayout(blockstype(arraytype)))