Skip to content

fix: mixed dense/sparse add!, scale! and inner #75

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lkdvos merged 1 commit into from
Sep 10, 2026
Merged

fix: mixed dense/sparse add!, scale! and inner #75

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 0 additions & 42 deletions src/linalg/linalg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,48 +9,6 @@ Base.:(*)(α::Number, t::AbstractBlockTensorMap) = scale(t, α)
Base.:(/)(t::AbstractBlockTensorMap, α::Number) = scale(t, inv(α))
Base.:(\)(α::Number, t::AbstractBlockTensorMap) = scale(t, inv(α))

function LinearAlgebra.axpy!(α::Number, t1::BlockTensorMap, t2::BlockTensorMap)
space(t1) == space(t2) || throw(SpaceMismatch())
for (i, v) in nonzero_pairs(t1)
t2[i] = axpy!(α, v, t2[i])
end
return t2
end

function LinearAlgebra.axpby!(α::Number, t1::BlockTensorMap, β::Number, t2::BlockTensorMap)
space(t1) == space(t2) || throw(SpaceMismatch())
rmul!(t2, β)
for (i, v) in nonzero_pairs(t1)
t2[i] = axpy!(α, v, t2[i])
end
return t2
end

function LinearAlgebra.dot(t1::BlockTensorMap, t2::BlockTensorMap)
size(t1) == size(t2) || throw(DimensionMismatch("dot arguments have different size"))

s = zero(promote_type(scalartype(t1), scalartype(t2)))
if nonzero_length(t1) >= nonzero_length(t2)
@inbounds for (I, v) in nonzero_pairs(t1)
s += dot(v, t2[I])
end
else
@inbounds for (I, v) in nonzero_pairs(t2)
s += dot(t1[I], v)
end
end
return s
end

function LinearAlgebra.mul!(C::BlockTensorMap, α::Number, A::BlockTensorMap)
space(C) == space(A) || throw(SpaceMismatch())
SparseArrayKit._zero!(parent(C))
for (i, v) in nonzero_pairs(A)
C[i] = mul!(C[i], α, v)
end
return C
end

for TA in (:AbstractBlockTensorMap, :(TK.DiagonalTensorMap), :(TK.AdjointTensorMap), :TensorMap),
TB in (:AbstractBlockTensorMap, :(TK.DiagonalTensorMap), :(TK.AdjointTensorMap), :TensorMap)
(TA === :AbstractBlockTensorMap || TB === :AbstractBlockTensorMap) || continue
Expand Down
69 changes: 19 additions & 50 deletions src/tensors/vectorinterface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,30 +10,13 @@ function VI.scale!(t::AbstractBlockTensorMap, α::Number)
return t
end

function VI.scale!(ty::BlockTensorMap, tx::BlockTensorMap, α::Number)
function VI.scale!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number)
space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))"))
scale!(parent(ty), parent(tx), α)
return ty
end
function VI.scale!(ty::SparseBlockTensorMap, tx::SparseBlockTensorMap, α::Number)
space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))"))
y_notin_x = setdiff(nonzero_keys(ty), nonzero_keys(tx))
x_notin_y = setdiff(nonzero_keys(tx), nonzero_keys(ty))
inboth = intersect(nonzero_keys(ty), nonzero_keys(tx))

# remove elements that are not in tx
for k in y_notin_x
delete!(ty.data, k)
end
# in-place scale elements that are in both
for k in inboth
ty[k] = scale!(ty[k], tx[k], α)
# entries of ty that are structurally zero in tx have to be zeroed out
issparse(tx) && zerovector!(ty)
for (I, v) in nonzero_pairs(tx)
ty[I] = scale!!(ty[I], v, α)
end
# new scale for elements in x that are not in y
for k in x_notin_y
ty[k] = scale(tx[k], α)
end

return ty
end

Expand Down Expand Up @@ -61,41 +44,27 @@ function VI.add(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Numb
return add!(scale!(tdst, ty, β), tx, α)
end

function VI.add!(ty::BlockTensorMap, tx::BlockTensorMap, α::Number, β::Number)
function VI.add!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number, β::Number)
space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))"))
add!(parent(ty), parent(tx), α, β)
return ty
end
function VI.add!(ty::SparseBlockTensorMap, tx::SparseBlockTensorMap, α::Number, β::Number)
space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))"))
y_notin_x = setdiff(nonzero_keys(ty), nonzero_keys(tx))
x_notin_y = setdiff(nonzero_keys(tx), nonzero_keys(ty))
inboth = intersect(nonzero_keys(ty), nonzero_keys(tx))

for k in y_notin_x
ty[k] = scale!!(ty[k], β)
end
for k in x_notin_y
ty[k] = scale(tx[k], α)
end
for k in inboth
ty[k] = add!!(ty[k], tx[k], α, β)
isone(β) || scale!(ty, β)
for (I, v) in nonzero_pairs(tx)
ty[I] = add!!(ty[I], v, α, One())
end

return ty
end

# inner
# -----
function VI.inner(x::BlockTensorMap, y::BlockTensorMap)
space(y) == space(x) || throw(SpaceMismatch())
return inner(parent(x), parent(y))
end
function VI.inner(x::SparseBlockTensorMap, y::SparseBlockTensorMap)
space(x) == space(y) || throw(SpaceMismatch())
both_nonzero = intersect(nonzero_keys(x), nonzero_keys(y))
function VI.inner(x::AbstractBlockTensorMap, y::AbstractBlockTensorMap)
space(x) == space(y) || throw(SpaceMismatch("$(space(x)) ≠ $(space(y))"))
T = VI.promote_inner(x, y)
return sum(both_nonzero; init = zero(T)) do k
inner(x[k], y[k])
# only entries that are nonzero in both contribute
ks = if issparse(x) && issparse(y)
intersect(nonzero_keys(x), nonzero_keys(y))
else
nonzero_keys(issparse(y) ? y : x)
end
return sum(ks; init = zero(T)) do I
inner(x[I], y[I])
end
end
57 changes: 47 additions & 10 deletions test/linalg/vectorinterface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,17 @@ Vtr = (

V = Vtr

@testset "VectorInterface $(issparse ? "SparseBlockTensorMap" : "BlockTensorMap")" for issparse in
(
false, true,
storagename(sparse) = sparse ? "SparseBlockTensorMap" : "BlockTensorMap"
makestorage(sparse, W) = sparse ? sprand(Float64, W, 0.5) : rand(W)

@testset "VectorInterface $(storagename(sparse_y)) ← $(storagename(sparse_x))" for (
sparse_y, sparse_x,
) in (
(false, false), (true, true), (false, true), (true, false),
)
if issparse
t = sprand(Float64, *(V[1:3]...) ← *(V[4], V[5]), 0.5)
t′ = sprand(Float64, *(V[1:3]...) ← *(V[4], V[5]), 0.5)
else
t = rand(*(V[1:3]...) ← *(V[4], V[5]))
t′ = rand(*(V[1:3]...) ← *(V[4], V[5]))
end
W = *(V[1:3]...) ← *(V[4], V[5])
t = makestorage(sparse_y, W)
t′ = makestorage(sparse_x, W)

@testset "scalartype" begin
@test Float64 === @inferred scalartype(t)
Expand Down Expand Up @@ -132,6 +132,15 @@ V = Vtr
@test @inferred inner(t″, t‴) ≈ conj(α) * β * inner(t, t′)
end

@testset "mixed-storage in-place" begin
α, β = rand(2)
@test @inferred scale!(zerovector(t), t′, α) ≈ scale(t′, α)
@test @inferred axpy!(α, t′, deepcopy(t)) ≈ add(t, t′, α)
@test @inferred axpby!(α, t′, β, deepcopy(t)) ≈ add(t, t′, α, β)
@test @inferred mul!(zerovector(t), α, t′) ≈ scale(t′, α)
@test @inferred dot(t, t′) ≈ inner(t, t′)
end

@testset "general linalg" begin
α, β = rand(ComplexF64, 2)
@test (α * α) * t ≈ α * (α * t)
Expand All @@ -143,3 +152,31 @@ V = Vtr
@test inner(t, t′) ≈ conj(inner(t′, t))
end
end

# regression test for #74: heterogeneous summands give interior zero-length coupled blocks
@testset "mixed storage with interior zero blocks" begin
S = Vect[U1Irrep]
P = S(0 => 1, 1 => 1, 2 => 1)
Vu = SumSpace(S(0 => 1), S(1 => 2), S(-1 => 2))
W = (Vu ⊗ SumSpace(P)) ← SumSpace(P)

tsp = SparseBlockTensorMap{TensorMap{Float64, S, 2, 1, Vector{Float64}}}(undef, W)
for (I, Wi) in zip(CartesianIndices(size(tsp)), BlockTensorKit.eachspace(tsp))
dim(Wi) == 0 || (tsp[I] = randn(Float64, Wi))
end
tdn = BlockTensorMap(tsp)
@test convert(TensorMap, tdn) ≈ convert(TensorMap, tsp)

α, β = rand(2)
ref = add(tdn, tdn, α, β)
for (y, x) in ((tdn, tsp), (tsp, tdn))
@test add!(deepcopy(y), x, α, β) ≈ ref
@test add(y, x, α, β) ≈ ref
@test y + x ≈ add(tdn, tdn)
@test norm(y - x) ≈ 0 atol = 1.0e-12
@test scale!(zerovector(y), x, α) ≈ scale(tdn, α)
@test axpy!(α, x, deepcopy(y)) ≈ add(tdn, tdn, α)
@test inner(y, x) ≈ inner(tdn, tdn)
@test dot(y, x) ≈ inner(tdn, tdn)
end
end
Loading