Merge branch 'main' into non-abelian_fusion

ITensor · Jun 18, 2024 · 2193a93 · 2193a93
2 parents b47fded + 7b2ada9
commit 2193a93
Show file tree

Hide file tree

Showing 12 changed files with 338 additions and 218 deletions.
diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml
@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.29"
+version = "0.3.31"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -42,8 +42,8 @@ TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
-NDTensorsAMDGPUExt = ["AMDGPU","GPUArraysCore"]
-NDTensorsCUDAExt = ["CUDA","GPUArraysCore"]
+NDTensorsAMDGPUExt = ["AMDGPU", "GPUArraysCore"]
+NDTensorsCUDAExt = ["CUDA", "GPUArraysCore"]
 NDTensorsGPUArraysCoreExt = "GPUArraysCore"
 NDTensorsHDF5Ext = "HDF5"
 NDTensorsMetalExt = ["GPUArraysCore", "Metal"]
@@ -59,7 +59,6 @@ ArrayLayouts = "1.4"
 BlockArrays = "1.1"
 CUDA = "5"
 Compat = "4.9"
-cuTENSOR = "2"
 Dictionaries = "0.4"
 EllipsisNotation = "1.8"
 FillArrays = "1"
@@ -81,11 +80,12 @@ SparseArrays = "1.6"
 SplitApplyCombine = "1.2.2"
 StaticArrays = "0.12, 1.0"
 Strided = "2"
-StridedViews = "0.2.2"
+StridedViews = "0.2.2, 0.3"
 TBLIS = "0.2"
 TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
 VectorInterface = "0.4.2"
+cuTENSOR = "2"
 julia = "1.6"
 
 [extras]

diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl
@@ -215,7 +215,7 @@ end
 
 function blockrange(
   axis::AbstractUnitRange,
-  r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
+  r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
 )
   return map(b -> Block(b), blocks(r))
 end
@@ -271,7 +271,7 @@ end
 function blockindices(
   a::AbstractUnitRange,
   b::Block,
-  r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
+  r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}},
 )
   # TODO: Change to iterate over `BlockRange(r)`
   # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404

diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl
@@ -27,7 +27,3 @@ end
 function Base.view(a::BlockSparseArrayLike{<:Any,1}, index::Block{1})
   return blocksparse_view(a, index)
 end
-
-function Base.view(a::BlockSparseArrayLike, indices::BlockIndexRange)
-  return view(view(a, block(indices)), indices.indices...)
-end
diff --git a/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl
@@ -120,19 +120,18 @@ blocktype(a::BlockSparseArrayLike) = eltype(blocks(a))
 blocktype(arraytype::Type{<:BlockSparseArrayLike}) = eltype(blockstype(arraytype))
 
 using ArrayLayouts: ArrayLayouts
-## function Base.getindex(a::BlockSparseArrayLike{<:Any,N}, I::Vararg{Int,N}) where {N}
-##   return ArrayLayouts.layout_getindex(a, I...)
-## end
 function Base.getindex(a::BlockSparseArrayLike{<:Any,N}, I::CartesianIndices{N}) where {N}
   return ArrayLayouts.layout_getindex(a, I)
 end
 function Base.getindex(
-  a::BlockSparseArrayLike{<:Any,N}, I::Vararg{AbstractUnitRange,N}
+  a::BlockSparseArrayLike{<:Any,N}, I::Vararg{AbstractUnitRange{<:Integer},N}
 ) where {N}
   return ArrayLayouts.layout_getindex(a, I...)
 end
 # TODO: Define `AnyBlockSparseMatrix`.
-function Base.getindex(a::BlockSparseArrayLike{<:Any,2}, I::Vararg{AbstractUnitRange,2})
+function Base.getindex(
+  a::BlockSparseArrayLike{<:Any,2}, I::Vararg{AbstractUnitRange{<:Integer},2}
+)
   return ArrayLayouts.layout_getindex(a, I...)
 end
 
@@ -199,7 +198,7 @@ end
 
 # Needed by `BlockArrays` matrix multiplication interface
 function Base.similar(
-  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractUnitRange}}
+  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
@@ -210,53 +209,45 @@ end
 # Delete once we drop support for older versions of Julia.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
-  axes::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}},
-)
-  return similar(arraytype, eltype(arraytype), axes)
-end
-
-# Needed by `BlockArrays` matrix multiplication interface
-# Fixes ambiguity error with `BlockArrays.jl`.
-function Base.similar(
-  arraytype::Type{<:BlockSparseArrayLike},
-  axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
+  axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
 
-# Needed by `BlockArrays` matrix multiplication interface
-# Fixes ambiguity error with `BlockArrays.jl`.
+# Fixes ambiguity error with `BlockArrays`.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
-  axes::Tuple{
-    AbstractBlockedUnitRange,AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}
-  },
+  axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
 
-# Needed by `BlockArrays` matrix multiplication interface
-# Fixes ambiguity error with `BlockArrays.jl`.
+# Fixes ambiguity error with `BlockArrays`.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
   axes::Tuple{
-    AbstractUnitRange{Int},AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}
+    AbstractUnitRange{<:Integer},
+    AbstractBlockedUnitRange{<:Integer},
+    Vararg{AbstractUnitRange{<:Integer}},
   },
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
 
 # Needed for disambiguation
 function Base.similar(
-  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractBlockedUnitRange}}
+  arraytype::Type{<:BlockSparseArrayLike},
+  axes::Tuple{Vararg{AbstractBlockedUnitRange{<:Integer}}},
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
 
 # Needed by `BlockArrays` matrix multiplication interface
 # TODO: Define a `blocksparse_similar` function.
 function Base.similar(
-  arraytype::Type{<:BlockSparseArrayLike}, elt::Type, axes::Tuple{Vararg{AbstractUnitRange}}
+  arraytype::Type{<:BlockSparseArrayLike},
+  elt::Type,
+  axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}},
 )
   # TODO: Make generic for GPU, maybe using `blocktype`.
   # TODO: For non-block axes this should output `Array`.
@@ -265,7 +256,7 @@ end
 
 # TODO: Define a `blocksparse_similar` function.
 function Base.similar(
-  a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{AbstractUnitRange}}
+  a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}
 )
   # TODO: Make generic for GPU, maybe using `blocktype`.
   # TODO: For non-block axes this should output `Array`.
@@ -277,7 +268,9 @@ end
 function Base.similar(
   a::BlockSparseArrayLike,
   elt::Type,
-  axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractBlockedUnitRange}},
+  axes::Tuple{
+    AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}}
+  },
 )
   # TODO: Make generic for GPU, maybe using `blocktype`.
   # TODO: For non-block axes this should output `Array`.
@@ -289,13 +282,37 @@ end
 function Base.similar(
   a::BlockSparseArrayLike,
   elt::Type,
-  axes::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}},
+  axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
   # TODO: Make generic for GPU, maybe using `blocktype`.
   # TODO: For non-block axes this should output `Array`.
   return BlockSparseArray{elt}(undef, axes)
 end
 
+# Fixes ambiguity error with `BlockArrays`.
+function Base.similar(
+  a::BlockSparseArrayLike,
+  elt::Type,
+  axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
+)
+  # TODO: Make generic for GPU, maybe using `blocktype`.
+  # TODO: For non-block axes this should output `Array`.
+  return BlockSparseArray{elt}(undef, axes)
+end
+
+# Fixes ambiguity errors with BlockArrays.
+function Base.similar(
+  a::BlockSparseArrayLike,
+  elt::Type,
+  axes::Tuple{
+    AbstractUnitRange{<:Integer},
+    AbstractBlockedUnitRange{<:Integer},
+    Vararg{AbstractUnitRange{<:Integer}},
+  },
+)
+  return BlockSparseArray{elt}(undef, axes)
+end
+
 # TODO: Define a `blocksparse_similar` function.
 # Fixes ambiguity error with `StaticArrays`.
 function Base.similar(

diff --git a/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
@@ -102,7 +102,24 @@ function blocksparse_fill!(a::AbstractArray, value)
     return a
   end
   for b in BlockRange(a)
-    a[b] .= value
+    # We can't use:
+    # ```julia
+    # a[b] .= value
+    # ```
+    # since that would lead to a stack overflow,
+    # because broadcasting calls `fill!`.
+
+    # TODO: Ideally we would use:
+    # ```julia
+    # @view!(a[b]) .= value
+    # ```
+    # but that doesn't work on `SubArray` right now.
+
+    # This line is needed to instantiate blocks
+    # that aren't instantiated yet. Maybe
+    # we can make this work without this line?
+    blocks(a)[Int.(Tuple(b))...] = blocks(a)[Int.(Tuple(b))...]
+    blocks(a)[Int.(Tuple(b))...] .= value
   end
   return a
 end
@@ -268,6 +285,10 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) w
 end
 function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N}
   parent_blocks = view(blocks(parent(a.array)), axes(a)...)
+  # TODO: The following line is required to instantiate
+  # uninstantiated blocks, maybe use `@view!` instead,
+  # or some other code pattern.
+  parent_blocks[I...] = parent_blocks[I...]
   return parent_blocks[I...][blockindices(parent(a.array), Block(I), a.array.indices)...] =
     value
 end