Merge branch 'main' into show

Author: lkdvos

.github/workflows/IntegrationTest.yml

@@ -0,0 +1,27 @@
+name: "IntegrationTest"
+
+on:
+  push:
+    branches:
+      - 'main'
+    tags: '*'
+    paths:
+      - 'Project.toml'
+  pull_request:
+    paths:
+      - 'Project.toml'
+
+jobs:
+  integration-test:
+    name: "IntegrationTest"
+    strategy:
+       matrix:
+         repo:
+           - 'ITensor/ITensorBase.jl'
+           - 'ITensor/NamedDimsArrays.jl'
+           - 'ITensor/QuantumOperatorDefinitions.jl'
+           - 'ITensor/TensorAlgebra.jl'
+    uses: "ITensor/ITensorActions/.github/workflows/IntegrationTest.yml@main"
+    with:
+      localregistry: "https://github.com/ITensor/ITensorRegistry.git"
+      repo: "${{ matrix.repo }}"

.github/workflows/Registrator.yml

@@ -9,9 +9,16 @@ on:
     branches:
       - 'master'
       - 'main'
+
+permissions:
+  contents: write
+  pull-requests: write
+
 jobs:
   Register:
     if: github.event_name == 'workflow_dispatch' || github.event.pull_request.merged == true
     uses: "ITensor/ITensorActions/.github/workflows/Registrator.yml@main"
     with:
       localregistry: ITensor/ITensorRegistry
+    secrets:
+      REGISTRATOR_KEY: ${{ secrets.REGISTRATOR_KEY }}

Project.toml

@@ -1,7 +1,8 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.2.18"
+version = "0.2.26"
+
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -45,7 +46,7 @@ MacroTools = "0.5.13"
 MapBroadcast = "0.1.5"
 SparseArraysBase = "0.2.10"
 SplitApplyCombine = "1.2.3"
-TensorAlgebra = "0.1.0"
+TensorAlgebra = "0.1.0, 0.2"
 Test = "1.10"
 TypeParameterAccessors = "0.2.0, 0.3"
 julia = "1.10"

ext/BlockSparseArraysGradedUnitRangesExt/BlockSparseArraysGradedUnitRangesExt.jl

@@ -2,7 +2,7 @@ module BlockSparseArraysGradedUnitRangesExt
 
 using BlockSparseArrays: AnyAbstractBlockSparseArray, BlockSparseArray, blocktype
 using GradedUnitRanges: AbstractGradedUnitRange
-using TypeParameterAccessors: set_eltype, set_ndims, unwrap_array_type
+using TypeParameterAccessors: similartype, unwrap_array_type
 
 # A block spare array similar to the input (dense) array.
 # TODO: Make `BlockSparseArrays.blocksparse_similar` more general and use that,
@@ -15,9 +15,7 @@ function similar_blocksparse(
   # TODO: Probably need to unwrap the type of `a` in certain cases
   # to make a proper block type.
   return BlockSparseArray{
-    elt,
-    length(axes),
-    set_eltype(set_ndims(unwrap_array_type(blocktype(a)), length(axes)), elt),
+    elt,length(axes),similartype(unwrap_array_type(blocktype(a)), elt, axes)
   }(
     axes
   )
@@ -52,6 +50,13 @@ function Base.similar(
 end
 
 # Fix ambiguity error with `BlockSparseArrays.jl`.
+function Base.similar(
+  a::AnyAbstractBlockSparseArray,
+  elt::Type,
+  axes::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}},
+)
+  return similar_blocksparse(a, elt, axes)
+end
 function Base.similar(
   a::AnyAbstractBlockSparseArray,
   elt::Type,

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysGradedUnitRangesExt.jl

@@ -94,13 +94,15 @@ function TensorAlgebra.splitdims(
     groupreducewhile(tensor_product, split_axes, ndims(a); init=OneToOne()) do i, axis
       return length(axis) ≤ length(axes(a, i))
     end
-  blockperms = invblockperm.(blocksortperm.(axes_prod))
+  blockperms = blocksortperm.(axes_prod)
+  sorted_axes = map((r, I) -> only(axes(r[I])), axes_prod, blockperms)
+
   # TODO: This is doing extra copies of the blocks,
   # use `@view a[axes_prod...]` instead.
   # That will require implementing some reindexing logic
   # for this combination of slicing.
-  a_unblocked = a[axes_prod...]
-  a_blockpermed = a_unblocked[blockperms...]
+  a_unblocked = a[sorted_axes...]
+  a_blockpermed = a_unblocked[invblockperm.(blockperms)...]
   return splitdims(BlockReshapeFusion(), a_blockpermed, split_axes...)
 end
 

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl

@@ -78,41 +78,52 @@ function Base.setindex!(
   return a
 end
 
-using TypeParameterAccessors: unspecify_type_parameters
-function show_typeof_blocksparse(io::IO, a::AbstractBlockSparseArray)
-  Base.show(io, unspecify_type_parameters(typeof(a)))
-  print(io, '{')
-  show(io, eltype(a))
-  print(io, ", ")
-  show(io, ndims(a))
-  print(io, ", ")
-  show(io, blocktype(a))
-  print(io, ", …")
-  print(io, '}')
-  return nothing
+# Copy of `Base.dims2string` defined in `show.jl`.
+function dims_to_string(d)
+  isempty(d) && return "0-dimensional"
+  length(d) == 1 && return "$(d[1])-element"
+  return join(map(string, d), '×')
 end
 
-# Copied from `BlockArrays.jl`.
-block2string(b, s) = string(join(map(string, b), '×'), "-blocked ", Base.dims2string(s))
+# Copy of `BlockArrays.block2string` from `BlockArrays.jl`.
+block_to_string(b, s) = string(join(map(string, b), '×'), "-blocked ", dims_to_string(s))
 
-function summary_blocksparse(io::IO, a::AbstractArray)
-  print(io, block2string(blocksize(a), size(a)))
-  print(io, ' ')
-  show_typeof_blocksparse(io, a)
-  return nothing
+using TypeParameterAccessors: type_parameters, unspecify_type_parameters
+function concretetype_to_string_truncated(type::Type; param_truncation_length=typemax(Int))
+  isconcretetype(type) || throw(ArgumentError("Type must be concrete."))
+  alias = Base.make_typealias(type)
+  base_type, params = if isnothing(alias)
+    unspecify_type_parameters(type), type_parameters(type)
+  else
+    base_type_globalref, params_svec = alias
+    base_type_globalref.name, params_svec
+  end
+  str = string(base_type)
+  if isempty(params)
+    return str
+  end
+  str *= '{'
+  param_strings = map(params) do param
+    param_string = string(param)
+    if length(param_string) > param_truncation_length
+      return "…"
+    end
+    return param_string
+  end
+  str *= join(param_strings, ", ")
+  str *= '}'
+  return str
 end
 
 function Base.summary(io::IO, a::AbstractBlockSparseArray)
-  summary_blocksparse(io, a)
+  print(io, block_to_string(blocksize(a), size(a)))
+  print(io, ' ')
+  print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length=40))
   return nothing
 end
 
 function Base.showarg(io::IO, a::AbstractBlockSparseArray, toplevel::Bool)
-  if toplevel
-    show_typeof_blocksparse(io, a)
-  else
-    print(io, "::")
-    show_typeof_blocksparse(io, a)
-  end
+  !toplevel && print(io, "::")
+  print(io, concretetype_to_string_truncated(typeof(a); param_truncation_length=40))
   return nothing
 end

src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -208,31 +208,3 @@ TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blocktype)) =
 function TypeParameterAccessors.position(::Type{BlockSparseArray}, ::typeof(blockstype))
   return Position(4)
 end
-
-# TODO: Make this generic to `AbstractBlockSparseVector` using
-# TypeParameterAccessors.jl, for example using:
-# `set_ndims(unspecify_type_parameters(typeof(a)), 1)`.
-function show_typeof_blocksparse(io::IO, a::BlockSparseVector)
-  print(io, "BlockSparseVector")
-  print(io, '{')
-  show(io, eltype(a))
-  print(io, ", ")
-  show(io, blocktype(a))
-  print(io, ", …")
-  print(io, '}')
-  return nothing
-end
-
-# TODO: Make this generic to `AbstractBlockSparseMatrix` using
-# TypeParameterAccessors.jl, for example using:
-# `set_ndims(unspecify_type_parameters(typeof(a)), 2)`.
-function show_typeof_blocksparse(io::IO, a::BlockSparseMatrix)
-  print(io, "BlockSparseMatrix")
-  print(io, '{')
-  show(io, eltype(a))
-  print(io, ", ")
-  show(io, blocktype(a))
-  print(io, ", …")
-  print(io, '}')
-  return nothing
-end

src/blocksparsearray/blocksparsearray.jl

@@ -31,10 +31,10 @@ function Broadcast.BroadcastStyle(
   return DefaultArrayStyle{N}()
 end
 
-function Base.similar(bc::Broadcasted{<:BlockSparseArrayStyle}, elt::Type)
-  # TODO: Make sure this handles GPU arrays properly.
+function Base.similar(bc::Broadcasted{<:BlockSparseArrayStyle}, elt::Type, ax)
+  # TODO: Make this more generic, base it off sure this handles GPU arrays properly.
   m = Mapped(bc)
-  return similar(first(m.args), elt, combine_axes(axes.(m.args)...))
+  return similar(first(m.args), elt, ax)
 end
 
 # Catches cases like `dest .= value` or `dest .= value1 .+ value2`.

src/blocksparsearrayinterface/broadcast.jl

@@ -337,6 +337,24 @@ arrayts = (Array, JLArray)
     @test blockstoredlength(a) == 1
     @test storedlength(a) == 2 * 4
 
+    # Test similar on broadcasted expressions.
+    a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4])))
+    bc = Broadcast.broadcasted(+, a, a)
+    a′ = similar(bc, Float32)
+    @test a′ isa BlockSparseArray{Float32}
+    @test blocktype(a′) <: arrayt{Float32,2}
+    @test axes(a) == (blockedrange([2, 3]), blockedrange([3, 4]))
+
+    # Test similar on broadcasted expressions with axes specified.
+    a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4])))
+    bc = Broadcast.broadcasted(+, a, a)
+    a′ = similar(
+      bc, Float32, (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2]))
+    )
+    @test a′ isa BlockSparseArray{Float32}
+    @test blocktype(a′) <: arrayt{Float32,3}
+    @test axes(a′) == (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2]))
+
     a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4])))
     @views for b in [Block(1, 2), Block(2, 1)]
       a[b] = dev(randn(elt, size(a[b])))
@@ -1090,34 +1108,30 @@ arrayts = (Array, JLArray)
     arrayt_elt = arrayt{elt,3}
 
     a = BlockSparseVector{elt,arrayt{elt,1}}([2, 2])
+    res = sprint(summary, a)
+    function ref_vec(elt, arrayt, prefix="")
+      return "2-blocked 4-element $(prefix)BlockSparseVector{$(elt), $(arrayt), …, …}"
+    end
     # Either option is possible depending on namespacing.
-    @test (
-      sprint(summary, a) ==
-      "2-blocked 4-element BlockSparseVector{$(elt), $(vectort_elt), …}"
-    ) || (
-      sprint(summary, a) ==
-      "2-blocked 4-element BlockSparseArrays.BlockSparseVector{$(elt), $(vectort_elt), …}"
-    )
+    @test (res == ref_vec(elt, vectort_elt)) ||
+      (res == ref_vec(elt, vectort_elt, "BlockSparseArrays."))
 
     a = BlockSparseMatrix{elt,arrayt{elt,2}}([2, 2], [2, 2])
+    res = sprint(summary, a)
+    function ref_mat(elt, arrayt, prefix="")
+      return "2×2-blocked 4×4 $(prefix)BlockSparseMatrix{$(elt), $(arrayt), …, …}"
+    end
     # Either option is possible depending on namespacing.
-    @test (
-      sprint(summary, a) == "2×2-blocked 4×4 BlockSparseMatrix{$(elt), $(matrixt_elt), …}"
-    ) || (
-      sprint(summary, a) ==
-      "2×2-blocked 4×4 BlockSparseArrays.BlockSparseMatrix{$(elt), $(matrixt_elt), …}"
-    )
+    @test (res == ref_mat(elt, matrixt_elt)) ||
+      (res == ref_mat(elt, matrixt_elt, "BlockSparseArrays."))
 
     a = BlockSparseArray{elt,3,arrayt{elt,3}}([2, 2], [2, 2], [2, 2])
-
-    # Either option is possible depending on namespacing.
-    @test (
-      sprint(summary, a) ==
-      "2×2×2-blocked 4×4×4 BlockSparseArray{$(elt), 3, $(arrayt_elt), …}"
-    ) || (
-      sprint(summary, a) ==
-      "2×2×2-blocked 4×4×4 BlockSparseArrays.BlockSparseArray{$(elt), 3, $(arrayt_elt), …}"
-    )
+    res = sprint(summary, a)
+    function ref_arr(elt, arrayt, prefix="")
+      return "2×2×2-blocked 4×4×4 $(prefix)BlockSparseArray{$(elt), 3, $(arrayt), …, …}"
+    end
+    @test (res == ref_arr(elt, arrayt_elt)) ||
+      (res == ref_arr(elt, arrayt_elt, "BlockSparseArrays."))
 
     if elt === Float64
       # Not testing other element types since they change the