From 3e33ad887876ea2cda2b2486759e4aa03af98f85 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:04:47 +0200
Subject: [PATCH 01/16] Modernize the benchmark harness

Give benchmarks a dedicated environment and use BenchmarkTools JSON serialization instead of the undeclared JLD/FileIO stack. Also repair report generation on current Julia.
---
 .gitignore                   |  3 ++-
 benchmark/Project.toml       |  8 ++++++++
 benchmark/README.md          | 23 +++++++++++++++++++++++
 benchmark/generate_report.jl |  6 +++++-
 benchmark/runbenchmarks.jl   | 21 +++++++++------------
 5 files changed, 47 insertions(+), 14 deletions(-)
 create mode 100644 benchmark/Project.toml
 create mode 100644 benchmark/README.md

diff --git a/.gitignore b/.gitignore
index ca28dd4d..3d39b6dc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,7 +4,8 @@
 docs/build/
 docs/site/
 *.jld
-benchmark/*.md
+benchmark/*.json
+benchmark/results_*.md
 src/.DS_Store
 Manifest.toml
 Manifest-v*.*.toml
diff --git a/benchmark/Project.toml b/benchmark/Project.toml
new file mode 100644
index 00000000..682e927a
--- /dev/null
+++ b/benchmark/Project.toml
@@ -0,0 +1,8 @@
+[deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[compat]
+BenchmarkTools = "1"
+julia = "1.10"
diff --git a/benchmark/README.md b/benchmark/README.md
new file mode 100644
index 00000000..7cdff9e0
--- /dev/null
+++ b/benchmark/README.md
@@ -0,0 +1,23 @@
+# Benchmarks
+
+Set up the benchmark environment from the repository root:
+
+```julia
+julia --project=benchmark -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate()'
+```
+
+Run and save the suite under a descriptive name:
+
+```julia
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'run_benchmarks("main")'
+```
+
+Generate a standalone report or compare two saved runs:
+
+```julia
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'generate_report("main")'
+julia --project=benchmark -L benchmark/runbenchmarks.jl -e 'generate_report("main", "branch")'
+```
+
+Benchmark parameters and results are stored as ignored JSON files under
+`benchmark/`.
diff --git a/benchmark/generate_report.jl b/benchmark/generate_report.jl
index be976b93..5786f4a6 100644
--- a/benchmark/generate_report.jl
+++ b/benchmark/generate_report.jl
@@ -40,7 +40,11 @@ function printreport(io::IO, results; iscomparisonjob::Bool = false)
     return nothing
 end
 
-idrepr(id) = (str = repr(id); str[searchindex(str, '['):end])
+function idrepr(id)
+    str = repr(id)
+    firstbracket = findfirst(==('['), str)
+    return isnothing(firstbracket) ? str : str[firstbracket:end]
+end
 intpercent(p) = string(ceil(Int, p * 100), "%")
 resultrow(ids, t::BenchmarkTools.Trial) = resultrow(ids, minimum(t))
 
diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index bfdaa202..a2b53b23 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -1,7 +1,5 @@
 using BlockArrays
 using BenchmarkTools
-using FileIO
-using JLD
 
 include("generate_report.jl")
 
@@ -32,29 +30,28 @@ end
 
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
-    paramspath = joinpath(dirname(@__FILE__), "params.jld")
+    paramspath = joinpath(@__DIR__, "params.json")
     if !isfile(paramspath)
         println("Tuning benchmarks...")
         tune!(SUITE, verbose=true)
-        JLD.save(paramspath, "SUITE", params(SUITE))
+        BenchmarkTools.save(paramspath, params(SUITE))
     end
-    loadparams!(SUITE, JLD.load(paramspath, "SUITE"), :evals, :samples)
+    loadparams!(SUITE, only(BenchmarkTools.load(paramspath)), :evals, :samples)
     results = run(SUITE[tagfilter], verbose = true, seconds = 2)
-    JLD.save(joinpath(dirname(@__FILE__), name * ".jld"), "results", results)
+    BenchmarkTools.save(joinpath(@__DIR__, name * ".json"), results)
 end
 
 function generate_report(v1, v2)
-    v1_res = load(joinpath(dirname(@__FILE__), v1 * ".jld"), "results")
-    v2_res = load(joinpath(dirname(@__FILE__), v2 * ".jld"), "results")
-    open(joinpath(dirname(@__FILE__), "results_$(v1)_$(v2).md"), "w") do f
+    v1_res = only(BenchmarkTools.load(joinpath(@__DIR__, v1 * ".json")))
+    v2_res = only(BenchmarkTools.load(joinpath(@__DIR__, v2 * ".json")))
+    open(joinpath(@__DIR__, "results_$(v1)_$(v2).md"), "w") do f
         printreport(f, judge(minimum(v1_res), minimum(v2_res)); iscomparisonjob = true)
     end
 end
 
 function generate_report(v1)
-    v1_res = load(joinpath(dirname(@__FILE__), v1 * ".jld"), "results")
-    open(joinpath(dirname(@__FILE__), "results_$(v1).md"), "w") do f
+    v1_res = only(BenchmarkTools.load(joinpath(@__DIR__, v1 * ".json")))
+    open(joinpath(@__DIR__, "results_$(v1).md"), "w") do f
         printreport(f, minimum(v1_res); iscomparisonjob = false)
     end
 end
-

From 1f23e2f9107797df691ca283ab9bfd870717a418 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:07:13 +0200
Subject: [PATCH 02/16] Make BlockKron axes allocation-free

Represent uniform Kronecker block lengths with Fill instead of allocating and accumulating dense vectors. Preserve implicit singleton dimensions for mixed vector/matrix products.

MWE: benchmark metadata/BlockKron/axes

Vector: 3.213 us, 160.12 KiB, 6 allocs -> 2.041 ns, 0 allocs
Matrix: 752.941 ns, 30.25 KiB, 12 allocs -> 2.292 ns, 0 allocs
---
 benchmark/runbenchmarks.jl |  8 ++++++++
 src/blockproduct.jl        | 14 +++++---------
 test/test_blockproduct.jl  |  7 +++++++
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index a2b53b23..929ce572 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -8,6 +8,7 @@ const SUITE = BenchmarkGroup()
 g = addgroup!(SUITE, "indexing")
 # g_block = addgroup!(SUITE, "blockindexing")
 g_size = addgroup!(SUITE, "size")
+g_metadata = addgroup!(SUITE, "metadata")
 
 for n = (5,)
     for BT in (BlockArray, BlockedArray)
@@ -28,6 +29,13 @@ for n = (5,)
     end
 end
 
+blockkron_vector = Ref(BlockKron(1:10_000, 1:4, 1:3))
+blockkron_matrix = Ref(BlockKron(
+    reshape(1:800_000, 1_000, 800), reshape(1:20, 4, 5), reshape(1:6, 3, 2),
+))
+g_metadata["BlockKron", "axes", "vector"] = @benchmarkable axes($blockkron_vector[])
+g_metadata["BlockKron", "axes", "matrix"] = @benchmarkable axes($blockkron_matrix[])
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockproduct.jl b/src/blockproduct.jl
index 19504822..64089851 100644
--- a/src/blockproduct.jl
+++ b/src/blockproduct.jl
@@ -51,15 +51,11 @@ size(K::BlockKron, j::Int) = prod(size.(K.args, j))
 size(a::BlockKron{<:Any,1}) = (size(a,1),)
 size(a::BlockKron{<:Any,2}) = (size(a,1), size(a,2))
 
-function axes(K::BlockKron{<:Any,1})
-    A,B = K.args
-    (blockedrange(fill(prod(size.(tail(K.args),1)), size(K.args[1],1))),)
-end
-
-function axes(K::BlockKron{<:Any,2})
-    A,B = K.args
-    blockedrange.((fill(prod(size.(tail(K.args),1)), size(K.args[1],1)),
-                   fill(prod(size.(tail(K.args),2)), size(K.args[1],2))))
+function axes(K::BlockKron{<:Any,N}) where N
+    ntuple(Val(N)) do dim
+        blocklength = prod(size.(tail(K.args), dim))
+        blockedrange(Fill(blocklength, size(K.args[1], dim)))
+    end
 end
 
 kron_getindex((A,)::Tuple{AbstractVector}, k::Integer) = A[k]
diff --git a/test/test_blockproduct.jl b/test/test_blockproduct.jl
index 13203f11..7c7ab55c 100644
--- a/test/test_blockproduct.jl
+++ b/test/test_blockproduct.jl
@@ -116,6 +116,8 @@ using BlockArrays, Test
         c = 6:8
         k̄ = BlockKron(a,b,c)
         @test k̄ == blockkron(a,b,c) == kron(a,b,c)
+        @test @inferred(axes(k̄)) == (blockedrange(fill(length(b) * length(c), length(a))),)
+        @test blocklasts(axes(k̄, 1)) isa FirstStepRangeLen
         @test k̄[Block(1)][Block(1)] == a[1]*b[1]*c
         @test k̄[Block(1)][Block(2)] == a[1]*b[2]*c
         @test k̄[Block(2)][Block(3)] == a[2]*b[3]*c
@@ -133,6 +135,11 @@ using BlockArrays, Test
         C = randn(2,5)
         K̄ = BlockKron(A,B,C)
         @test K̄ ≈ blockkron(A,B,C) ≈ kron(A,B,C)
+        @test @inferred(axes(K̄)) == (
+            blockedrange(fill(size(B, 1) * size(C, 1), size(A, 1))),
+            blockedrange(fill(size(B, 2) * size(C, 2), size(A, 2))),
+        )
+        @test all(ax -> blocklasts(ax) isa FirstStepRangeLen, axes(K̄))
         @test K̄[Block(1,1)][Block(1,1)] ≈ A[1,1]*B[1,1]*C
         @test K̄[Block(2,3)][Block(3,4)] ≈ A[2,3]*B[3,4]*C
 

From c9cd60fe889c8abf9ec20fd1226a50eb8fcade52 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:11:04 +0200
Subject: [PATCH 03/16] Merge sorted block boundaries linearly

Specialize sortedunion for strided integer vectors so broadcast axis combination does one linear pass instead of hashing and sorting already-sorted boundaries.

MWE: benchmark metadata/sortedunion/vectors (10k + 10k boundaries)

Before: 215.334 us, 1.33 MiB, 53 allocs
After: 20.458 us, 160.06 KiB, 3 allocs
---
 benchmark/runbenchmarks.jl  |  5 +++++
 src/blockbroadcast.jl       | 25 ++++++++++++++++++++++++-
 test/test_blockbroadcast.jl |  7 +++++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 929ce572..0f164b20 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -36,6 +36,11 @@ blockkron_matrix = Ref(BlockKron(
 g_metadata["BlockKron", "axes", "vector"] = @benchmarkable axes($blockkron_vector[])
 g_metadata["BlockKron", "axes", "matrix"] = @benchmarkable axes($blockkron_matrix[])
 
+blocklasts_a = collect(2:2:20_000)
+blocklasts_b = collect(3:3:30_000)
+g_metadata["sortedunion", "vectors"] =
+    @benchmarkable BlockArrays.sortedunion($blocklasts_a, $blocklasts_b)
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockbroadcast.jl b/src/blockbroadcast.jl
index c1180268..4f796a0f 100644
--- a/src/blockbroadcast.jl
+++ b/src/blockbroadcast.jl
@@ -33,10 +33,33 @@ BroadcastStyle(::BlockStyle{M}, ::BlockedStyle{N}) where {M,N} = BlockStyle(Val(
 BroadcastStyle(::BlockedStyle{M}, ::BlockStyle{N}) where {M,N} = BlockStyle(Val(max(M,N)))
 
 
-# sortedunion can assume inputs are already sorted so this could be improved
 maybeinplacesort!(v::StridedVector) = sort!(v)
 maybeinplacesort!(v) = sort(v)
 sortedunion(a,b) = maybeinplacesort!(union(a,b))
+function sortedunion(a::StridedVector{<:Integer}, b::StridedVector{<:Integer})
+    T = promote_type(eltype(a), eltype(b))
+    result = Vector{T}(undef, length(a) + length(b))
+    ia, ib = firstindex(a), firstindex(b)
+    lasta, lastb = lastindex(a), lastindex(b)
+    nresult = 0
+
+    @inbounds while ia <= lasta || ib <= lastb
+        value = if ib > lastb || (ia <= lasta && !isless(b[ib], a[ia]))
+            value = a[ia]
+            ia += 1
+            value
+        else
+            value = b[ib]
+            ib += 1
+            value
+        end
+        if iszero(nresult) || !isequal(result[nresult], value)
+            nresult += 1
+            result[nresult] = value
+        end
+    end
+    return resize!(result, nresult)
+end
 sortedunion(a::Base.OneTo, b::Base.OneTo) = Base.OneTo(max(last(a),last(b)))
 sortedunion(a::AbstractUnitRange, b::AbstractUnitRange) = min(first(a),first(b)):max(last(a),last(b))
 combine_blockaxes(a, b) = _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
diff --git a/test/test_blockbroadcast.jl b/test/test_blockbroadcast.jl
index cbfdd97c..469f624b 100644
--- a/test/test_blockbroadcast.jl
+++ b/test/test_blockbroadcast.jl
@@ -129,6 +129,13 @@ using StaticArrays
         @test blocksize(A+B) == (5,3)
     end
 
+    @testset "sorted block boundary union" begin
+        @test BlockArrays.sortedunion([1, 3, 3, 7], [2, 3, 8]) == [1, 2, 3, 7, 8]
+        @test BlockArrays.sortedunion(Int[], Int[]) == Int[]
+        @test BlockArrays.sortedunion(Int32[1, 3], Int64[2, 3]) == [1, 2, 3]
+        @test eltype(BlockArrays.sortedunion(Int32[1], Int64[2])) == Int64
+    end
+
     @testset "UnitRange" begin
         n = 3
         x = mortar([1:4n, 1:n])

From 90b126b22a6a76549015e85a60a9f0c8fad0169f Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:12:12 +0200
Subject: [PATCH 04/16] Make khatri_rao type-stable

Build result blocks with a typed product map and take non-copying views of the inputs. This removes the nested Any vectors and temporary copies of every input block.

MWE: benchmark product/khatri_rao/10x10/2x2 blocks

Before: 14.458 us, 63.78 KiB, 702 allocs, return type Any
After: 2.588 us, 23.44 KiB, 212 allocs, concrete BlockMatrix
---
 benchmark/runbenchmarks.jl |  7 +++++++
 src/blockproduct.jl        | 23 +++++++----------------
 test/test_blockproduct.jl  |  3 ++-
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 0f164b20..33389498 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -9,6 +9,7 @@ g = addgroup!(SUITE, "indexing")
 # g_block = addgroup!(SUITE, "blockindexing")
 g_size = addgroup!(SUITE, "size")
 g_metadata = addgroup!(SUITE, "metadata")
+g_product = addgroup!(SUITE, "product")
 
 for n = (5,)
     for BT in (BlockArray, BlockedArray)
@@ -41,6 +42,12 @@ blocklasts_b = collect(3:3:30_000)
 g_metadata["sortedunion", "vectors"] =
     @benchmarkable BlockArrays.sortedunion($blocklasts_a, $blocklasts_b)
 
+khatri_block_sizes = fill(2, 10)
+khatri_a = BlockArray(randn(20, 20), khatri_block_sizes, khatri_block_sizes)
+khatri_b = BlockArray(randn(20, 20), khatri_block_sizes, khatri_block_sizes)
+g_product["khatri_rao", "10x10", "2x2 blocks"] =
+    @benchmarkable khatri_rao($khatri_a, $khatri_b)
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockproduct.jl b/src/blockproduct.jl
index 64089851..baa982b5 100644
--- a/src/blockproduct.jl
+++ b/src/blockproduct.jl
@@ -6,23 +6,14 @@ References
 * Khatri, C. G., and Rao, C. Radhakrishna (1968) Solutions to Some Functional Equations and Their Applications to Characterization of Probability Distributions. Sankhya: Indian J. Statistics, Series A 30, 167–180.
 """
 function khatri_rao(A::AbstractBlockMatrix, B::AbstractBlockMatrix)
-    #
-    Ablksize = blocksize(A)
-    Bblksize = blocksize(B)
-
-    @assert Ablksize == Bblksize "A and B must have the same blocksize"
-
-    kblk = []
-    for iblk in blockaxes(A,1)
-        kblk_j = []
-        for _jblk in blockaxes(A,2)
-            Ablk = A[iblk, _jblk]
-            Bblk = B[iblk, _jblk]
-            push!(kblk_j, kron(Ablk, Bblk))
-        end
-        push!(kblk, tuple(kblk_j...))
+    @assert blocksize(A) == blocksize(B) "A and B must have the same blocksize"
+
+    product = Iterators.product(blockaxes(A)...)
+    result_blocks = map(product) do block_index
+        K, J = block_index
+        kron(view(A, K, J), view(B, K, J))
     end
-    mortar(kblk...)
+    return mortar(result_blocks)
 end
 
 function khatri_rao(A::AbstractMatrix, B::AbstractMatrix)
diff --git a/test/test_blockproduct.jl b/test/test_blockproduct.jl
index 7c7ab55c..0dd0286c 100644
--- a/test/test_blockproduct.jl
+++ b/test/test_blockproduct.jl
@@ -27,10 +27,11 @@ using BlockArrays, Test
         A = BlockArray(ones(m, n), mi, ni)
         B = BlockArray(ones(p, q), pi, qi)
 
-        AB = khatri_rao(A, B)
+        AB = @inferred khatri_rao(A, B)
 
         @test blocksize(AB) == blocksize(A)
         @test blocksize(AB) == blocksize(B)
+        @test khatri_rao(BlockedArray(A), BlockedArray(B)) == AB
 
         #Test: Size of resulting blocks
         for i in blockaxes(AB,1)

From 0c1f10b38f5cb3688555634674b0b94e56e7ddc3 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:14:40 +0200
Subject: [PATCH 05/16] Delegate BlockedArray broadcasts to parent storage

Unwrap pure BlockedStyle broadcasts for in-place assignment and reuse identical block axes instead of rebuilding their boundaries. Mismatched partitions continue through the existing axis-combination path.

MWE: benchmark broadcast/in-place/BlockedArray/matching (128x128, 32x32 blocks)

Before: 29.041 us, 2.25 KiB, 8 allocs
After: 2.940 us, 0 bytes, 0 allocs
---
 benchmark/runbenchmarks.jl  |  8 ++++++++
 src/blockbroadcast.jl       | 11 +++++++++--
 test/test_blockbroadcast.jl | 13 +++++++++++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 33389498..71e792d4 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -10,6 +10,7 @@ g = addgroup!(SUITE, "indexing")
 g_size = addgroup!(SUITE, "size")
 g_metadata = addgroup!(SUITE, "metadata")
 g_product = addgroup!(SUITE, "product")
+g_broadcast = addgroup!(SUITE, "broadcast")
 
 for n = (5,)
     for BT in (BlockArray, BlockedArray)
@@ -48,6 +49,13 @@ khatri_b = BlockArray(randn(20, 20), khatri_block_sizes, khatri_block_sizes)
 g_product["khatri_rao", "10x10", "2x2 blocks"] =
     @benchmarkable khatri_rao($khatri_a, $khatri_b)
 
+broadcast_block_sizes = fill(4, 32)
+broadcast_blocked_a = BlockedArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_blocked_b = BlockedArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_blocked_dest = similar(broadcast_blocked_a)
+g_broadcast["in-place", "BlockedArray", "matching"] =
+    @benchmarkable $broadcast_blocked_dest .= $broadcast_blocked_a .+ $broadcast_blocked_b
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockbroadcast.jl b/src/blockbroadcast.jl
index 4f796a0f..b54244f9 100644
--- a/src/blockbroadcast.jl
+++ b/src/blockbroadcast.jl
@@ -62,8 +62,8 @@ function sortedunion(a::StridedVector{<:Integer}, b::StridedVector{<:Integer})
 end
 sortedunion(a::Base.OneTo, b::Base.OneTo) = Base.OneTo(max(last(a),last(b)))
 sortedunion(a::AbstractUnitRange, b::AbstractUnitRange) = min(first(a),first(b)):max(last(a),last(b))
-combine_blockaxes(a, b) = _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
-combine_blockaxes(a::BlockedOneTo, b::BlockedOneTo) = BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+combine_blockaxes(a, b) = blockisequal(a, b) ? a : _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
+combine_blockaxes(a::BlockedOneTo, b::BlockedOneTo) = blockisequal(a, b) ? a : BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
 
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b::AbstractBlockedUnitRange) = length(b) == 1 ? a : combine_blockaxes(a, b)
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b) = length(b) == 1 ? a : combine_blockaxes(a, b)
@@ -237,6 +237,13 @@ _removeblocks(a::Adjoint) = _removeblocks(parent(a))'
 _removeblocks(a::Transpose) = transpose(_removeblocks(parent(a)))
 _removeblocks(a::SubArray{<:Any,N,<:BlockedArray}) where N = view(_removeblocks(parent(a)), map(_removeblocks, parentindices(a))...)
 _removeblocks(a) = a
+
+function copyto!(dest::BlockedArray{<:Any,N},
+        bc::Broadcasted{BlockedStyle{N},<:Any,<:Any,Args}) where {N,Args<:Tuple}
+    copyto!(parent(dest), _removeblocks(bc))
+    return dest
+end
+
 copy(bc::Broadcasted{BlockedStyle{N}}) where N = BlockedArray(Broadcast.materialize(_removeblocks(bc)), axes(bc))
 
 for op in (:+, :-, :*)
diff --git a/test/test_blockbroadcast.jl b/test/test_blockbroadcast.jl
index 469f624b..4468e811 100644
--- a/test/test_blockbroadcast.jl
+++ b/test/test_blockbroadcast.jl
@@ -81,6 +81,19 @@ using StaticArrays
             @test dest ≈ x + 2y
         end
 
+        @testset "matrix in-place broadcast" begin
+            x = BlockedMatrix(randn(6, 6), [2, 4], [3, 3])
+            y = BlockedMatrix(randn(6, 6), [1, 2, 3], [2, 4])
+            dest = similar(x)
+            @test (dest .= x .+ 2 .* y) === dest
+            @test parent(dest) ≈ parent(x) .+ 2 .* parent(y)
+            @test axes(x .+ x) === axes(x)
+
+            expected = parent(x) .+ parent(y)
+            x .+= y
+            @test parent(x) ≈ expected
+        end
+
         @testset "0-dim nested in-place broadcast" begin
             x = BlockedArray(randn(()))
             y = BlockedArray(randn(()))

From b95da7f4437e290542f082b17d9256ec6248bcb5 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:16:58 +0200
Subject: [PATCH 06/16] Broadcast matching arrays by N-dimensional blocks

Generalize the vector-only fast path to any dimensionality when every array argument has exactly the destination's block axes. Dimension expansion and mismatched partitions retain the generic sub-block iterator.

MWE: benchmark broadcast/in-place/BlockArray/matching (128x128, 32x32 blocks)

Before: 44.833 us, 0 allocs
After: 13.417 us, 0 allocs
---
 benchmark/runbenchmarks.jl  |  6 ++++++
 src/blockbroadcast.jl       | 27 +++++++++++----------------
 test/test_blockbroadcast.jl | 13 +++++++++++++
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 71e792d4..27ac2f5e 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -56,6 +56,12 @@ broadcast_blocked_dest = similar(broadcast_blocked_a)
 g_broadcast["in-place", "BlockedArray", "matching"] =
     @benchmarkable $broadcast_blocked_dest .= $broadcast_blocked_a .+ $broadcast_blocked_b
 
+broadcast_block_a = BlockArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_block_b = BlockArray(randn(128, 128), broadcast_block_sizes, broadcast_block_sizes)
+broadcast_block_dest = similar(broadcast_block_a)
+g_broadcast["in-place", "BlockArray", "matching"] =
+    @benchmarkable $broadcast_block_dest .= $broadcast_block_a .+ $broadcast_block_b
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockbroadcast.jl b/src/blockbroadcast.jl
index b54244f9..f89ddd07 100644
--- a/src/blockbroadcast.jl
+++ b/src/blockbroadcast.jl
@@ -199,31 +199,26 @@ copyto!(dest::AbstractArray,
     _generic_blockbroadcast_copyto!(dest, bc)
 
 # type-stable version of _bview.(args, K)
-__bview(args::Tuple{}, K) = ()
-__bview(args::Tuple, K) = tuple(_bview(args[1],K), __bview(tail(args), K)...)
+__bview(args::Tuple{}, K...) = ()
+__bview(args::Tuple, K...) = tuple(_bview(args[1], K...), __bview(tail(args), K...)...)
 
-function _fast_blockbradcast_copyto!(dest, bc)
-    @inbounds for K in blockaxes(bc)[1]
-        broadcast!(bc.f, view(dest,K), __bview(bc.args, K)...)
+function _fast_blockbroadcast_copyto!(dest, bc)
+    @inbounds for K in Iterators.product(blockaxes(bc)...)
+        broadcast!(bc.f, view(dest, K...), __bview(bc.args, K...)...)
     end
     dest
 end
 
-_hasscalarlikevec() = false
-_hasscalarlikevec(a, b...) = _hasscalarlikevec(b...)
-_hasscalarlikevec(a::AbstractVector, b...) = size(a,1) == 1 || _hasscalarlikevec(b...)
+blockisequalorscalar(ax::Tuple, ::Number) = true
+blockisequalorscalar(ax::Tuple, a) = blockisequal(ax, axes(a))
 
-blockisequalorscalar(ax, ::Number) = true
-blockisequalorscalar(ax, a) = blockisequal(ax, Base.axes1(a))
-
-function copyto!(dest::AbstractVector,
-        bc::Broadcasted{<:AbstractBlockStyle{1}, <:Any, <:Any, Args}) where {Args <: Tuple}
-    _hasscalarlikevec(bc.args...) && return _generic_blockbroadcast_copyto!(dest, bc)
-    ax = axes(dest,1)
+function copyto!(dest::AbstractArray{<:Any,N},
+        bc::Broadcasted{<:AbstractBlockStyle{N},<:Any,<:Any,Args}) where {N,Args<:Tuple}
+    ax = axes(dest)
     for a in bc.args
         blockisequalorscalar(ax, a) || return _generic_blockbroadcast_copyto!(dest, bc)
     end
-    return _fast_blockbradcast_copyto!(dest, bc)
+    return _fast_blockbroadcast_copyto!(dest, bc)
 end
 @inline function Broadcast.instantiate(bc::Broadcasted{Style}) where {Style <:BlockStyle}
     bcf = Broadcast.instantiate(Broadcast.flatten(Broadcasted{Nothing}(bc.f, bc.args, bc.axes)))
diff --git a/test/test_blockbroadcast.jl b/test/test_blockbroadcast.jl
index 4468e811..12364aa2 100644
--- a/test/test_blockbroadcast.jl
+++ b/test/test_blockbroadcast.jl
@@ -28,6 +28,15 @@ using StaticArrays
         @test axes(A + A) == axes(A .+ A) == axes(A)
         @test axes(A .+ 1) == axes(A)
 
+        dest = similar(A)
+        @test (dest .= A .+ 2 .* A) === dest
+        @test Matrix(dest) ≈ Matrix(A) .+ 2 .* Matrix(A)
+
+        A3 = BlockArray(randn(4, 4, 4), [2, 2], [1, 3], [3, 1])
+        dest3 = similar(A3)
+        @test (dest3 .= A3 .+ A3) === dest3
+        @test Array(dest3) ≈ 2 .* Array(A3)
+
         @testset "mismatched ndims" begin
             u = BlockArray(randn(5), [2,3])
             dest = zeros(size(u)..., 1)
@@ -140,6 +149,10 @@ using StaticArrays
         B = BlockArray(randn(6,6), fill(2,3), fill(3,2))
 
         @test blocksize(A+B) == (5,3)
+
+        dest = similar(A)
+        @test (dest .= A .+ B) === dest
+        @test Matrix(dest) ≈ Matrix(A) .+ Matrix(B)
     end
 
     @testset "sorted block boundary union" begin

From 893b59c48ff25c94164917a7a7c80ddfc1599a3f Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:19:17 +0200
Subject: [PATCH 07/16] Reduce BlockArray sums per stored block

Specialize whole-array sum and mapped sum without changing the existing dims path. Apply init only once at the outer reduction and retain Base's generic identity handling when there are no stored blocks.

MWE: benchmark reduction/{sum,sum(abs2)}/BlockArray/matrix (128x128, 32x32 blocks)

sum: 90.375 us -> 3.479 us (26.0x)
sum(abs2): 91.000 us -> 4.101 us (22.2x)
Both remain allocation-free.
---
 benchmark/runbenchmarks.jl |  7 +++++++
 src/blockreduce.jl         |  8 ++++++++
 test/test_blockreduce.jl   | 12 +++++++++++-
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 27ac2f5e..3582c227 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -11,6 +11,7 @@ g_size = addgroup!(SUITE, "size")
 g_metadata = addgroup!(SUITE, "metadata")
 g_product = addgroup!(SUITE, "product")
 g_broadcast = addgroup!(SUITE, "broadcast")
+g_reduction = addgroup!(SUITE, "reduction")
 
 for n = (5,)
     for BT in (BlockArray, BlockedArray)
@@ -62,6 +63,12 @@ broadcast_block_dest = similar(broadcast_block_a)
 g_broadcast["in-place", "BlockArray", "matching"] =
     @benchmarkable $broadcast_block_dest .= $broadcast_block_a .+ $broadcast_block_b
 
+reduction_block_sizes = fill(4, 32)
+reduction_block_array = BlockArray(randn(128, 128), reduction_block_sizes, reduction_block_sizes)
+g_reduction["sum", "BlockArray", "matrix"] = @benchmarkable sum($reduction_block_array)
+g_reduction["sum(abs2)", "BlockArray", "matrix"] =
+    @benchmarkable sum(abs2, $reduction_block_array)
+
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
     paramspath = joinpath(@__DIR__, "params.json")
diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index f1c55363..b01a36fe 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -17,5 +17,13 @@ Base.mapfoldl(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
 Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
     mapreduce(f, op, B.blocks; kw...)
 
+Base.sum(B::BlockArray; dims=:, kw...) = sum(identity, B; dims, kw...)
+function Base.sum(f, B::BlockArray; dims=:, kw...)
+    if dims isa Colon && !isempty(B.blocks)
+        return mapreduce(block -> sum(f, block), Base.add_sum, B.blocks; kw...)
+    end
+    return invoke(sum, Tuple{Any,AbstractArray}, f, B; dims, kw...)
+end
+
 # support sum, need to return something analogous to Base.OneTo(1) but same type
 Base.reduced_index(::BR) where BR<:AbstractBlockedUnitRange = convert(BR, Base.OneTo(1))
diff --git a/test/test_blockreduce.jl b/test/test_blockreduce.jl
index 5af3afab..c794a295 100644
--- a/test/test_blockreduce.jl
+++ b/test/test_blockreduce.jl
@@ -21,11 +21,21 @@ end
 @testset "sum (#141)" begin
     data = reshape(collect(1:20), 4, 5)
     A = BlockArray(data, [1,3], [2,3])
-    @test sum(A) == sum(data)
+    @test @inferred(sum(A)) == sum(data)
+    @test @inferred(sum(abs2, A)) == sum(abs2, data)
+    @test sum(A; init=10) == sum(data; init=10)
     @test sum(A; dims=1) == sum(data; dims=1)
     @test sum(A; dims=2) == sum(data; dims=2)
     @test blockisequal(axes(A,2), axes(sum(A; dims=1),2))
     @test blockisequal(axes(A,1), axes(sum(A; dims=2),1))
+
+    smallints = BlockArray(fill(Int8(1), 4, 4), [2, 2], [1, 3])
+    @test sum(smallints) === sum(Matrix(smallints)) === 16
+
+    emptyblocks = BlockArray(zeros(0, 0), Int[], Int[])
+    @test sum(emptyblocks) === 0.0
+    @test sum(abs2, emptyblocks) === 0.0
+    @test sum(emptyblocks; init=10.0) === 10.0
 end
 
 end # module

From c75e9d108c7164e368be25c3f887fe6b85dbc255 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 12:21:00 +0200
Subject: [PATCH 08/16] Compute BlockArray 2-norm from block norms

Combine each stored block's stable 2-norm with hypot. Other p-norms and arrays with no stored blocks continue through the generic implementation.

MWE: benchmark reduction/norm/BlockArray/matrix (128x128, 32x32 blocks)

Before: 191.750 us, 48 bytes, 1 alloc
After: 18.666 us, 0 bytes, 0 allocs
---
 benchmark/Project.toml     |  1 +
 benchmark/runbenchmarks.jl |  2 ++
 src/blockreduce.jl         |  7 +++++++
 test/test_blocklinalg.jl   | 13 +++++++++++++
 4 files changed, 23 insertions(+)

diff --git a/benchmark/Project.toml b/benchmark/Project.toml
index 682e927a..f740db19 100644
--- a/benchmark/Project.toml
+++ b/benchmark/Project.toml
@@ -1,6 +1,7 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [compat]
diff --git a/benchmark/runbenchmarks.jl b/benchmark/runbenchmarks.jl
index 3582c227..55d3e3d7 100644
--- a/benchmark/runbenchmarks.jl
+++ b/benchmark/runbenchmarks.jl
@@ -1,5 +1,6 @@
 using BlockArrays
 using BenchmarkTools
+using LinearAlgebra
 
 include("generate_report.jl")
 
@@ -68,6 +69,7 @@ reduction_block_array = BlockArray(randn(128, 128), reduction_block_sizes, reduc
 g_reduction["sum", "BlockArray", "matrix"] = @benchmarkable sum($reduction_block_array)
 g_reduction["sum(abs2)", "BlockArray", "matrix"] =
     @benchmarkable sum(abs2, $reduction_block_array)
+g_reduction["norm", "BlockArray", "matrix"] = @benchmarkable norm($reduction_block_array)
 
 
 function run_benchmarks(name, tagfilter = @tagged ALL)
diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index b01a36fe..71652ce9 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -25,5 +25,12 @@ function Base.sum(f, B::BlockArray; dims=:, kw...)
     return invoke(sum, Tuple{Any,AbstractArray}, f, B; dims, kw...)
 end
 
+function LinearAlgebra.norm(B::BlockArray, p::Real=2)
+    if p == 2 && !isempty(B.blocks)
+        return mapreduce(norm, hypot, B.blocks)
+    end
+    return invoke(norm, Tuple{Any,Real}, B, p)
+end
+
 # support sum, need to return something analogous to Base.OneTo(1) but same type
 Base.reduced_index(::BR) where BR<:AbstractBlockedUnitRange = convert(BR, Base.OneTo(1))
diff --git a/test/test_blocklinalg.jl b/test/test_blocklinalg.jl
index d10f0d70..bc7156ca 100644
--- a/test/test_blocklinalg.jl
+++ b/test/test_blocklinalg.jl
@@ -35,6 +35,19 @@ bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
     @test a .^ 2 == 4
     end
 
+    @testset "norm" begin
+        data = randn(7, 9)
+        A = BlockArray(data, [1, 2, 4], [3, 1, 5])
+        @test @inferred(norm(A)) ≈ norm(data)
+        @test norm(A, 1) ≈ norm(data, 1)
+
+        extremes = BlockArray([1.0e300 1.0e-300], [1], [1, 1])
+        @test norm(extremes) == norm(Matrix(extremes)) == 1.0e300
+
+        emptyblocks = BlockArray(zeros(0, 0), Int[], Int[])
+        @test norm(emptyblocks) === 0.0
+    end
+
     @testset "BlockArray scalar * matrix" begin
         A = BlockArray{Float64}(randn(6,6), fill(2,3), 1:3)
         @test 2A == A*2 == 2Matrix(A)

From d1df6b59d76884996d00963b607aac7aba950bc2 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 16:35:59 +0200
Subject: [PATCH 09/16] Keep sorted boundary merging with axis helpers

There is no public Base vector merge that exploits sorted inputs. Co-locate BlockArrays' linear vector specialization with the existing lazy sortedunion axis methods, and document why the local implementation is needed.
---
 src/blockaxis.jl      | 28 ++++++++++++++++++++++++++++
 src/blockbroadcast.jl | 24 ------------------------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/src/blockaxis.jl b/src/blockaxis.jl
index 1acfd5d4..c71ae4dd 100644
--- a/src/blockaxis.jl
+++ b/src/blockaxis.jl
@@ -735,6 +735,34 @@ Base.BroadcastStyle(::Type{<:AbstractBlockedUnitRange{<:Any,R}}) where R = _broa
 ###
 
 const OneToCumsum{T<:Integer} = RangeCumsum{T,Base.OneTo{T}}
+
+# Base has no public vector merge that exploits sorted inputs, so merge block
+# boundaries directly instead of hashing and then sorting them with `union`.
+function sortedunion(a::StridedVector{<:Integer}, b::StridedVector{<:Integer})
+    T = promote_type(eltype(a), eltype(b))
+    result = Vector{T}(undef, length(a) + length(b))
+    ia, ib = firstindex(a), firstindex(b)
+    lasta, lastb = lastindex(a), lastindex(b)
+    nresult = 0
+
+    @inbounds while ia <= lasta || ib <= lastb
+        value = if ib > lastb || (ia <= lasta && !isless(b[ib], a[ia]))
+            value = a[ia]
+            ia += 1
+            value
+        else
+            value = b[ib]
+            ib += 1
+            value
+        end
+        if iszero(nresult) || !isequal(result[nresult], value)
+            nresult += 1
+            result[nresult] = value
+        end
+    end
+    return resize!(result, nresult)
+end
+
 sortedunion(a::OneToCumsum, ::OneToCumsum) = a
 function sortedunion(a::RangeCumsum{<:Any,<:AbstractRange}, b::RangeCumsum{<:Any,<:AbstractRange})
     @assert a == b
diff --git a/src/blockbroadcast.jl b/src/blockbroadcast.jl
index f89ddd07..60a44afa 100644
--- a/src/blockbroadcast.jl
+++ b/src/blockbroadcast.jl
@@ -36,30 +36,6 @@ BroadcastStyle(::BlockedStyle{M}, ::BlockStyle{N}) where {M,N} = BlockStyle(Val(
 maybeinplacesort!(v::StridedVector) = sort!(v)
 maybeinplacesort!(v) = sort(v)
 sortedunion(a,b) = maybeinplacesort!(union(a,b))
-function sortedunion(a::StridedVector{<:Integer}, b::StridedVector{<:Integer})
-    T = promote_type(eltype(a), eltype(b))
-    result = Vector{T}(undef, length(a) + length(b))
-    ia, ib = firstindex(a), firstindex(b)
-    lasta, lastb = lastindex(a), lastindex(b)
-    nresult = 0
-
-    @inbounds while ia <= lasta || ib <= lastb
-        value = if ib > lastb || (ia <= lasta && !isless(b[ib], a[ia]))
-            value = a[ia]
-            ia += 1
-            value
-        else
-            value = b[ib]
-            ib += 1
-            value
-        end
-        if iszero(nresult) || !isequal(result[nresult], value)
-            nresult += 1
-            result[nresult] = value
-        end
-    end
-    return resize!(result, nresult)
-end
 sortedunion(a::Base.OneTo, b::Base.OneTo) = Base.OneTo(max(last(a),last(b)))
 sortedunion(a::AbstractUnitRange, b::AbstractUnitRange) = min(first(a),first(b)):max(last(a),last(b))
 combine_blockaxes(a, b) = blockisequal(a, b) ? a : _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))

From dc8f97cb71303176b7d27da6a4bdc934bc518d85 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 16:38:56 +0200
Subject: [PATCH 10/16] Keep matching block-axis combination inferred

Only reuse an input axis when the matching and merged branches have the same concrete representation. Add inference tests for reusable, heterogeneous integer, and lazy block axes.
---
 src/blockbroadcast.jl       | 14 ++++++++++++--
 test/test_blockbroadcast.jl |  8 ++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/blockbroadcast.jl b/src/blockbroadcast.jl
index 60a44afa..bb74900d 100644
--- a/src/blockbroadcast.jl
+++ b/src/blockbroadcast.jl
@@ -38,8 +38,18 @@ maybeinplacesort!(v) = sort(v)
 sortedunion(a,b) = maybeinplacesort!(union(a,b))
 sortedunion(a::Base.OneTo, b::Base.OneTo) = Base.OneTo(max(last(a),last(b)))
 sortedunion(a::AbstractUnitRange, b::AbstractUnitRange) = min(first(a),first(b)):max(last(a),last(b))
-combine_blockaxes(a, b) = blockisequal(a, b) ? a : _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
-combine_blockaxes(a::BlockedOneTo, b::BlockedOneTo) = blockisequal(a, b) ? a : BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+combine_blockaxes(a, b) = _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
+combine_blockaxes(a::BlockedOneTo, b::BlockedOneTo) = BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+
+function combine_blockaxes(a::BlockedUnitRange{T,Vector{T}},
+                           b::BlockedUnitRange{T,Vector{T}}) where T
+    return blockisequal(a, b) ? a : _BlockedUnitRange(sortedunion(blocklasts(a), blocklasts(b)))
+end
+
+function combine_blockaxes(a::BlockedOneTo{T,CS}, b::BlockedOneTo{T,CS}) where
+                           {T<:Integer,CS<:Union{Vector{T},Base.OneTo{T}}}
+    return blockisequal(a, b) ? a : BlockedOneTo(sortedunion(blocklasts(a), blocklasts(b)))
+end
 
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b::AbstractBlockedUnitRange) = length(b) == 1 ? a : combine_blockaxes(a, b)
 Base.Broadcast.axistype(a::AbstractBlockedUnitRange, b) = length(b) == 1 ? a : combine_blockaxes(a, b)
diff --git a/test/test_blockbroadcast.jl b/test/test_blockbroadcast.jl
index 12364aa2..307ea258 100644
--- a/test/test_blockbroadcast.jl
+++ b/test/test_blockbroadcast.jl
@@ -202,10 +202,18 @@ using StaticArrays
     @testset "special axes" begin
         A = BlockArray(randn(6), Ones{Int}(6))
         B = BlockArray(randn(6), Ones{Int}(6))
+        @test (@inferred BlockArrays.combine_blockaxes(axes(A, 1), axes(B, 1))) === axes(A, 1)
         @test axes(A+B,1) === axes(A,1)
 
         C = BlockArray(randn(6), (BlockArrays._BlockedUnitRange(1,2:6),))
         @test axes(A+C,1) === BlockArrays._BlockedUnitRange(1,1:6)
+
+        a32 = blockedrange(Int32[1, 2, 3])
+        a64 = blockedrange(Int64[1, 2, 3])
+        @test blockisequal(@inferred(BlockArrays.combine_blockaxes(a32, a64)), a64)
+
+        lazy = blockedrange(Fill(4, 32))
+        @test blockisequal(@inferred(BlockArrays.combine_blockaxes(lazy, lazy)), lazy)
     end
 
     @testset "Views" begin

From a3918d86ef4e7960aea2f8117193c59d51605a26 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 16:39:39 +0200
Subject: [PATCH 11/16] Route plain BlockArray sums through mapreduce

Put the per-block identity/add_sum reduction at the mapreduce level and make sum(::BlockArray) delegate to it. Direct mapreduce calls now share the optimized path while dims and empty storage retain Base's generic behavior.
---
 src/blockreduce.jl       | 14 +++++++++++++-
 test/test_blockreduce.jl |  7 +++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index 71652ce9..98c88420 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -17,7 +17,19 @@ Base.mapfoldl(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
 Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
     mapreduce(f, op, B.blocks; kw...)
 
-Base.sum(B::BlockArray; dims=:, kw...) = sum(identity, B; dims, kw...)
+Base.sum(B::BlockArray; dims=:, kw...) = mapreduce(identity, Base.add_sum, B; dims, kw...)
+function Base.mapreduce(::typeof(identity), op::typeof(Base.add_sum), B::BlockArray;
+                        dims=:, kw...)
+    if dims isa Colon && !isempty(B.blocks)
+        return mapreduce(block -> mapreduce(identity, op, block), op, B.blocks; kw...)
+    end
+    return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, identity, op, B; dims, kw...)
+end
+Base.mapreduce(::typeof(identity), op::typeof(Base.add_sum), B::BlockVector;
+               dims=:, kw...) =
+    invoke(mapreduce, Tuple{typeof(identity),typeof(op),BlockArray},
+           identity, op, B; dims, kw...)
+
 function Base.sum(f, B::BlockArray; dims=:, kw...)
     if dims isa Colon && !isempty(B.blocks)
         return mapreduce(block -> sum(f, block), Base.add_sum, B.blocks; kw...)
diff --git a/test/test_blockreduce.jl b/test/test_blockreduce.jl
index c794a295..47996545 100644
--- a/test/test_blockreduce.jl
+++ b/test/test_blockreduce.jl
@@ -22,6 +22,9 @@ end
     data = reshape(collect(1:20), 4, 5)
     A = BlockArray(data, [1,3], [2,3])
     @test @inferred(sum(A)) == sum(data)
+    @test @inferred(mapreduce(identity, Base.add_sum, A)) == sum(data)
+    @test mapreduce(identity, Base.add_sum, A; init=10) == sum(data; init=10)
+    @test mapreduce(identity, Base.add_sum, A; dims=1) == sum(data; dims=1)
     @test @inferred(sum(abs2, A)) == sum(abs2, data)
     @test sum(A; init=10) == sum(data; init=10)
     @test sum(A; dims=1) == sum(data; dims=1)
@@ -36,6 +39,10 @@ end
     @test sum(emptyblocks) === 0.0
     @test sum(abs2, emptyblocks) === 0.0
     @test sum(emptyblocks; init=10.0) === 10.0
+
+    v = BlockArray(collect(1:6), [2, 4])
+    @test @inferred(sum(v)) == 21
+    @test @inferred(mapreduce(identity, Base.add_sum, v)) == 21
 end
 
 end # module

From 735875f2740bf8a8daad2f87a045386528d97960 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 16:44:21 +0200
Subject: [PATCH 12/16] Route mapped BlockArray sums through mapreduce

Generalize the add_sum mapreduce specialization to mapped reductions and make sum(f, ::BlockArray) delegate to it. Cover direct mapped mapreduce calls, init, dims fallback, and BlockVector dispatch.
---
 src/blockreduce.jl       | 22 ++++++++--------------
 test/test_blockreduce.jl |  5 +++++
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index 98c88420..744f55f5 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -18,24 +18,18 @@ Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
     mapreduce(f, op, B.blocks; kw...)
 
 Base.sum(B::BlockArray; dims=:, kw...) = mapreduce(identity, Base.add_sum, B; dims, kw...)
-function Base.mapreduce(::typeof(identity), op::typeof(Base.add_sum), B::BlockArray;
-                        dims=:, kw...)
+function Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockArray;
+                        dims=:, kw...) where F
     if dims isa Colon && !isempty(B.blocks)
-        return mapreduce(block -> mapreduce(identity, op, block), op, B.blocks; kw...)
+        return mapreduce(block -> mapreduce(f, op, block), op, B.blocks; kw...)
     end
-    return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, identity, op, B; dims, kw...)
+    return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, f, op, B; dims, kw...)
 end
-Base.mapreduce(::typeof(identity), op::typeof(Base.add_sum), B::BlockVector;
-               dims=:, kw...) =
-    invoke(mapreduce, Tuple{typeof(identity),typeof(op),BlockArray},
-           identity, op, B; dims, kw...)
+Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockVector;
+               dims=:, kw...) where F =
+    invoke(mapreduce, Tuple{F,typeof(op),BlockArray}, f, op, B; dims, kw...)
 
-function Base.sum(f, B::BlockArray; dims=:, kw...)
-    if dims isa Colon && !isempty(B.blocks)
-        return mapreduce(block -> sum(f, block), Base.add_sum, B.blocks; kw...)
-    end
-    return invoke(sum, Tuple{Any,AbstractArray}, f, B; dims, kw...)
-end
+Base.sum(f, B::BlockArray; dims=:, kw...) = mapreduce(f, Base.add_sum, B; dims, kw...)
 
 function LinearAlgebra.norm(B::BlockArray, p::Real=2)
     if p == 2 && !isempty(B.blocks)
diff --git a/test/test_blockreduce.jl b/test/test_blockreduce.jl
index 47996545..f79f232c 100644
--- a/test/test_blockreduce.jl
+++ b/test/test_blockreduce.jl
@@ -26,6 +26,9 @@ end
     @test mapreduce(identity, Base.add_sum, A; init=10) == sum(data; init=10)
     @test mapreduce(identity, Base.add_sum, A; dims=1) == sum(data; dims=1)
     @test @inferred(sum(abs2, A)) == sum(abs2, data)
+    @test @inferred(mapreduce(abs2, Base.add_sum, A)) == sum(abs2, data)
+    @test mapreduce(abs2, Base.add_sum, A; init=10) == sum(abs2, data; init=10)
+    @test mapreduce(abs2, Base.add_sum, A; dims=2) == sum(abs2, data; dims=2)
     @test sum(A; init=10) == sum(data; init=10)
     @test sum(A; dims=1) == sum(data; dims=1)
     @test sum(A; dims=2) == sum(data; dims=2)
@@ -43,6 +46,8 @@ end
     v = BlockArray(collect(1:6), [2, 4])
     @test @inferred(sum(v)) == 21
     @test @inferred(mapreduce(identity, Base.add_sum, v)) == 21
+    @test @inferred(sum(abs2, v)) == 91
+    @test @inferred(mapreduce(abs2, Base.add_sum, v)) == 91
 end
 
 end # module

From 8a7ba0d4a558e8429a79e0196dca38ae9509f329 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 16:45:38 +0200
Subject: [PATCH 13/16] Specialize BlockArray norm2

Extend LinearAlgebra's dedicated norm2 hook instead of norm itself. Generic norm dispatch now handles all other p values without an invoke, while the blockwise stable hypot reduction remains allocation-free.
---
 src/blockreduce.jl       | 8 +++-----
 test/test_blocklinalg.jl | 3 +++
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index 744f55f5..ff3a4fb6 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -31,11 +31,9 @@ Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockVector;
 
 Base.sum(f, B::BlockArray; dims=:, kw...) = mapreduce(f, Base.add_sum, B; dims, kw...)
 
-function LinearAlgebra.norm(B::BlockArray, p::Real=2)
-    if p == 2 && !isempty(B.blocks)
-        return mapreduce(norm, hypot, B.blocks)
-    end
-    return invoke(norm, Tuple{Any,Real}, B, p)
+function LinearAlgebra.norm2(B::BlockArray)
+    isempty(B.blocks) && return float(norm(zero(eltype(B))))
+    return mapreduce(norm, hypot, B.blocks)
 end
 
 # support sum, need to return something analogous to Base.OneTo(1) but same type
diff --git a/test/test_blocklinalg.jl b/test/test_blocklinalg.jl
index bc7156ca..ef4564dc 100644
--- a/test/test_blocklinalg.jl
+++ b/test/test_blocklinalg.jl
@@ -39,13 +39,16 @@ bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
         data = randn(7, 9)
         A = BlockArray(data, [1, 2, 4], [3, 1, 5])
         @test @inferred(norm(A)) ≈ norm(data)
+        @test @inferred(LinearAlgebra.norm2(A)) ≈ norm(data)
         @test norm(A, 1) ≈ norm(data, 1)
+        @test norm(A, 3) ≈ norm(data, 3)
 
         extremes = BlockArray([1.0e300 1.0e-300], [1], [1, 1])
         @test norm(extremes) == norm(Matrix(extremes)) == 1.0e300
 
         emptyblocks = BlockArray(zeros(0, 0), Int[], Int[])
         @test norm(emptyblocks) === 0.0
+        @test LinearAlgebra.norm2(emptyblocks) === 0.0
     end
 
     @testset "BlockArray scalar * matrix" begin

From d0cd82d0408f7f45e99c25a7167c05d5a6b9946e Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 18:22:58 +0200
Subject: [PATCH 14/16] Use Base sum delegation

Base already routes AbstractArray sums and mapped sums through mapreduce for whole-array reductions. Remove the redundant BlockArray methods and let the add_sum specialization remain the single extension point.
---
 src/blockreduce.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index ff3a4fb6..67eec6a1 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -17,7 +17,6 @@ Base.mapfoldl(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
 Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
     mapreduce(f, op, B.blocks; kw...)
 
-Base.sum(B::BlockArray; dims=:, kw...) = mapreduce(identity, Base.add_sum, B; dims, kw...)
 function Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockArray;
                         dims=:, kw...) where F
     if dims isa Colon && !isempty(B.blocks)
@@ -29,8 +28,6 @@ Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockVector;
                dims=:, kw...) where F =
     invoke(mapreduce, Tuple{F,typeof(op),BlockArray}, f, op, B; dims, kw...)
 
-Base.sum(f, B::BlockArray; dims=:, kw...) = mapreduce(f, Base.add_sum, B; dims, kw...)
-
 function LinearAlgebra.norm2(B::BlockArray)
     isempty(B.blocks) && return float(norm(zero(eltype(B))))
     return mapreduce(norm, hypot, B.blocks)

From 6bf445d5ec81e6e3af431c73c170d1c44e90b9b6 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 18:23:10 +0200
Subject: [PATCH 15/16] Skip empty blocks in mapped sums

A logically nonempty BlockArray may still store zero-length blocks. Exclude them from the outer per-block reduction so mapped sums do not require an additive identity for an empty block when the overall input is nonempty.
---
 src/blockreduce.jl       |  5 +++--
 test/test_blockreduce.jl | 10 ++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index 67eec6a1..8c09b688 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -19,8 +19,9 @@ Base.mapreduce(f::F, op::OP, B::BlockedArray; kw...) where {F, OP} =
 
 function Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockArray;
                         dims=:, kw...) where F
-    if dims isa Colon && !isempty(B.blocks)
-        return mapreduce(block -> mapreduce(f, op, block), op, B.blocks; kw...)
+    if dims isa Colon && !isempty(B)
+        nonemptyblocks = Iterators.filter(!isempty, B.blocks)
+        return mapreduce(block -> mapreduce(f, op, block), op, nonemptyblocks; kw...)
     end
     return invoke(mapreduce, Tuple{Any,Any,AbstractArray}, f, op, B; dims, kw...)
 end
diff --git a/test/test_blockreduce.jl b/test/test_blockreduce.jl
index f79f232c..56b1fd7a 100644
--- a/test/test_blockreduce.jl
+++ b/test/test_blockreduce.jl
@@ -2,6 +2,11 @@ module TestBlockReduce
 
 using BlockArrays, Test
 
+struct NoZero
+    value::Int
+end
+Base.:+(a::NoZero, b::NoZero) = NoZero(a.value + b.value)
+
 @testset "foldl" begin
     x = mortar([rand(3), rand(2)])
     @test foldl(push!, x; init = []) == collect(x)
@@ -43,6 +48,11 @@ end
     @test sum(abs2, emptyblocks) === 0.0
     @test sum(emptyblocks; init=10.0) === 10.0
 
+    zeroblocks = BlockArray(collect(1:6), [0, 2, 0, 4])
+    @test sum(zeroblocks) == sum(Vector(zeroblocks))
+    @test sum(abs2, zeroblocks) == sum(abs2, Vector(zeroblocks))
+    @test sum(NoZero, zeroblocks).value == sum(NoZero, Vector(zeroblocks)).value
+
     v = BlockArray(collect(1:6), [2, 4])
     @test @inferred(sum(v)) == 21
     @test @inferred(mapreduce(identity, Base.add_sum, v)) == 21

From aa4f109a5c493e4cfd8ebff6e5ac8ea4a096a804 Mon Sep 17 00:00:00 2001
From: Kristoffer Carlsson <kcarlsson89@gmail.com>
Date: Wed, 12 Aug 2026 18:23:20 +0200
Subject: [PATCH 16/16] Preserve nonfinite BlockArray norm semantics

Julia's generic 2-norm propagates NaN even when another element is infinite, whereas hypot(NaN, Inf) returns Inf. Use a combining function that propagates NaN and missing before applying hypot.
---
 src/blockreduce.jl       | 5 ++++-
 test/test_blocklinalg.jl | 6 ++++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/blockreduce.jl b/src/blockreduce.jl
index 8c09b688..873a662a 100644
--- a/src/blockreduce.jl
+++ b/src/blockreduce.jl
@@ -29,9 +29,12 @@ Base.mapreduce(f::F, op::typeof(Base.add_sum), B::BlockVector;
                dims=:, kw...) where F =
     invoke(mapreduce, Tuple{F,typeof(op),BlockArray}, f, op, B; dims, kw...)
 
+_norm_hypot(x, y) = (ismissing(x) || isnan(x)) ? x :
+                    (ismissing(y) || isnan(y)) ? y : hypot(x, y)
+
 function LinearAlgebra.norm2(B::BlockArray)
     isempty(B.blocks) && return float(norm(zero(eltype(B))))
-    return mapreduce(norm, hypot, B.blocks)
+    return mapreduce(norm, _norm_hypot, B.blocks)
 end
 
 # support sum, need to return something analogous to Base.OneTo(1) but same type
diff --git a/test/test_blocklinalg.jl b/test/test_blocklinalg.jl
index ef4564dc..23cbe060 100644
--- a/test/test_blocklinalg.jl
+++ b/test/test_blocklinalg.jl
@@ -46,6 +46,12 @@ bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)
         extremes = BlockArray([1.0e300 1.0e-300], [1], [1, 1])
         @test norm(extremes) == norm(Matrix(extremes)) == 1.0e300
 
+        nonfinite = BlockArray([NaN, Inf], [1, 1])
+        @test isnan(norm(nonfinite)) == isnan(norm(Vector(nonfinite)))
+
+        withmissing = BlockArray(Union{Missing,Float64}[1.0, missing], [1, 1])
+        @test norm(withmissing) === norm(Vector(withmissing)) === missing
+
         emptyblocks = BlockArray(zeros(0, 0), Int[], Int[])
         @test norm(emptyblocks) === 0.0
         @test LinearAlgebra.norm2(emptyblocks) === 0.0