BlockArraysExt

Author: mtfishman

Project.toml

@@ -4,12 +4,15 @@ authors = ["ITensor developers <support@itensor.org> and contributors"]
 version = "0.3.1"
 
 [weakdeps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [extensions]
+FunctionImplementationsBlockArraysExt = "BlockArrays"
 FunctionImplementationsLinearAlgebraExt = "LinearAlgebra"
 
 [compat]
+BlockArrays = "1.4"
 LinearAlgebra = "1.10"
 julia = "1.10"
 

ext/FunctionImplementationsBlockArraysExt/FunctionImplementationsBlockArraysExt.jl

@@ -0,0 +1,11 @@
+module FunctionImplementationsBlockArraysExt
+
+using BlockArrays: AbstractBlockedUnitRange, blockedrange, blocklengths
+using FunctionImplementations.Concatenate: Concatenate
+
+function Concatenate.cat_axis(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)
+    first(a1) == first(a2) == 1 || throw(ArgumentError("Concatenated axes must start at 1"))
+    return blockedrange([blocklengths(a1); blocklengths(a2)])
+end
+
+end

src/concatenate.jl

@@ -89,16 +89,19 @@ end
 
 function Base.similar(concat::Concatenated, ::Type{T}, ax) where {T}
     # Convert to a broadcasted to leverage its similar implementation.
-    bc = BC.Broadcasted(style(concat), nothing, ())
-    return similar(bc, T, ax)
+    bc = BC.Broadcasted(style(concat), identity, concat.args, ax)
+    return similar(bc, T)
 end
 
 function cat_axis(
         a1::AbstractUnitRange, a2::AbstractUnitRange, a_rest::AbstractUnitRange...
     )
     return cat_axis(cat_axis(a1, a2), a_rest...)
 end
-cat_axis(a1::AbstractUnitRange, a2::AbstractUnitRange) = Base.OneTo(length(a1) + length(a2))
+function cat_axis(a1::AbstractUnitRange, a2::AbstractUnitRange)
+    first(a1) == first(a2) == 1 || throw(ArgumentError("Concatenated axes must start at 1"))
+    return Base.OneTo(length(a1) + length(a2))
+end
 
 function cat_ndims(dims, as::AbstractArray...)
     return max(maximum(dims), maximum(ndims, as))

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 FunctionImplementations = "7c7cc465-9c6a-495f-bdd1-f42428e86d0c"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -14,6 +15,7 @@ FunctionImplementations = {path = ".."}
 [compat]
 Adapt = "4"
 Aqua = "0.8"
+BlockArrays = "1.4"
 FunctionImplementations = "0.3"
 JLArrays = "0.3"
 LinearAlgebra = "1.10"

test/test_blockarraysext.jl

@@ -0,0 +1,17 @@
+using BlockArrays: BlockArray, blockedrange, blockisequal
+using FunctionImplementations.Concatenate: concatenate
+using Test: @test, @testset
+
+@testset "BlockArraysExt" begin
+    a = BlockArray(randn(4, 4), [2, 2], [2, 2])
+    b = BlockArray(randn(4, 4), [2, 2], [2, 2])
+
+    concat = concatenate(1, a, b)
+    @test axes(concat) == (Base.OneTo(8), Base.OneTo(4))
+    @test blockisequal(axes(concat, 1), blockedrange([2, 2, 2, 2]))
+    @test blockisequal(axes(concat, 2), blockedrange([2, 2]))
+    @test size(concat) == (8, 4)
+    @test eltype(concat) ≡ Float64
+    @test copy(concat) == cat(a, b; dims = 1)
+    @test copy(concat) isa BlockArray{Float64, 2}
+end