Skip to content

Forward AD tests and cleanup for vector interface #436

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
kshyatt merged 2 commits into from
Jun 25, 2026
Merged

Forward AD tests and cleanup for vector interface #436

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
63afe12
Forward rules for vector interface
kshyatt May 19, 2026
3add6ac
Drop VI versions that don't have AD rules
kshyatt Jun 25, 2026
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
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,5 +64,5 @@ Strided = "2.6.1"
TensorKitSectors = "0.3.7"
TensorOperations = "5.5.2"
TupleTools = "1.5"
VectorInterface = "0.4.8, 0.5, 0.6"
VectorInterface = "0.6"
julia = "1.10"
2 changes: 1 addition & 1 deletion ext/TensorKitMooncakeExt/tangent.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ _field_symbol(t, ::Val{F}) where {F} = _field_symbol(t, F)

# frules
_frule_getfield_common(t_dt::Dual{<:DiagOrTensorMap}, field_sym::Symbol) =
Dual(getfield(primal(t), field_sym), field_sym === :data ? tangent(t).data : NoFData())
Dual(getfield(primal(t_dt), field_sym), field_sym === :data ? tangent(t_dt).data : NoTangent())

Mooncake.frule!!(::Dual{typeof(Mooncake.lgetfield)}, t_dt::Dual{<:DiagOrTensorMap}, f_df::Dual) =
_frule_getfield_common(t_dt, _field_symbol(primal(t_dt), primal(f_df)))
Expand Down
1 change: 1 addition & 0 deletions ext/TensorKitMooncakeExt/utility.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ Mooncake.tangent_type(::Type{<:HomSpace}) = Mooncake.NoTangent
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.has_shared_permute), AbstractTensorMap, Index2Tuple}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.select), HomSpace, Index2Tuple}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.flip), HomSpace, Any}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.adjoint), HomSpace}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.permute), HomSpace, Index2Tuple}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.braid), HomSpace, Index2Tuple, IndexTuple}
@zero_derivative DefaultCtx Tuple{typeof(TensorKit.compose), HomSpace, HomSpace}
Expand Down
93 changes: 4 additions & 89 deletions ext/TensorKitMooncakeExt/vectorinterface.jl
View file
Open in desktop
Comment thread
lkdvos marked this conversation as resolved.
Original file line number Diff line number Diff line change
@@ -1,89 +1,4 @@
@is_primitive DefaultCtx ReverseMode Tuple{typeof(scale!), AbstractTensorMap, Number}

function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractTensorMap}, α_Δα::CoDual{<:Number})
# prepare arguments
C, ΔC = arrayify(C_ΔC)
α = primal(α_Δα)

# primal call
C_cache = copy(C)
scale!(C, α)

function scale_pullback(::NoRData)
copy!(C, C_cache)
Δαr = _needs_tangent(α) ? TO.project_scalar(α, inner(C, ΔC)) : NoRData()
scale!(ΔC, conj(α))
return NoRData(), NoRData(), Δαr
end

return C_ΔC, scale_pullback
end

@is_primitive DefaultCtx ReverseMode Tuple{typeof(scale!), AbstractTensorMap, AbstractTensorMap, Number}

function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractTensorMap}, A_ΔA::CoDual{<:AbstractTensorMap}, α_Δα::CoDual{<:Number})
# prepare arguments
C, ΔC = arrayify(C_ΔC)
A, ΔA = arrayify(A_ΔA)
α = primal(α_Δα)

# primal call
C_cache = copy(C)
scale!(C, A, α)

function scale_pullback(::NoRData)
copy!(C, C_cache)
add!(ΔA, ΔC, conj(α))
Δαr = _needs_tangent(α) ? TO.project_scalar(α, inner(A, ΔC)) : NoRData()
zerovector!(ΔC)
return NoRData(), NoRData(), NoRData(), Δαr
end

return C_ΔC, scale_pullback
end

@is_primitive DefaultCtx ReverseMode Tuple{typeof(add!), AbstractTensorMap, AbstractTensorMap, Number, Number}

function Mooncake.rrule!!(::CoDual{typeof(add!)}, C_ΔC::CoDual{<:AbstractTensorMap}, A_ΔA::CoDual{<:AbstractTensorMap}, α_Δα::CoDual{<:Number}, β_Δβ::CoDual{<:Number})
# prepare arguments
C, ΔC = arrayify(C_ΔC)
A, ΔA = arrayify(A_ΔA)
α = primal(α_Δα)
β = primal(β_Δβ)

# primal call
C_cache = copy(C)
add!(C, A, α, β)

function add_pullback(::NoRData)
copy!(C, C_cache)

Δαr = _needs_tangent(α) ? TO.project_scalar(α, inner(A, ΔC)) : NoRData()
Δβr = _needs_tangent(β) ? TO.project_scalar(β, inner(C, ΔC)) : NoRData()
add!(ΔA, ΔC, conj(α))
scale!(ΔC, conj(β))

return NoRData(), NoRData(), NoRData(), Δαr, Δβr
end

return C_ΔC, add_pullback
end

@is_primitive DefaultCtx ReverseMode Tuple{typeof(inner), AbstractTensorMap, AbstractTensorMap}

function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual{<:AbstractTensorMap}, B_ΔB::CoDual{<:AbstractTensorMap})
# prepare arguments
A, ΔA = arrayify(A_ΔA)
B, ΔB = arrayify(B_ΔB)

# primal call
s = inner(A, B)

function inner_pullback(Δs)
add!(ΔA, B, conj(Δs))
add!(ΔB, A, Δs)
return NoRData(), NoRData(), NoRData()
end

return CoDual(s, NoFData()), inner_pullback
end
@is_primitive DefaultCtx Tuple{typeof(scale!), AbstractTensorMap, Number}
@is_primitive DefaultCtx Tuple{typeof(scale!), AbstractTensorMap, AbstractTensorMap, Number}
@is_primitive DefaultCtx Tuple{typeof(add!), AbstractTensorMap, AbstractTensorMap, Number, Number}
@is_primitive DefaultCtx Tuple{typeof(inner), AbstractTensorMap, AbstractTensorMap}
38 changes: 19 additions & 19 deletions test/mooncake/vectorinterface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
using Test, TestExtras
using TensorKit
using TensorOperations
using Mooncake
using VectorInterface, Mooncake
using Random


mode = Mooncake.ReverseMode
rng = Random.default_rng()

spacelist = ad_spacelist(fast_tests)
Expand All @@ -17,20 +15,22 @@ eltypes = (Float64, ComplexF64)

C = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
α = randn(T)
β = randn(T)

Mooncake.TestUtils.test_rule(rng, scale!, C, α; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, scale!, C', α; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, scale!, C, A, α; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, scale!, C', A', α; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, scale!, copy(C'), A', α; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, scale!, C', copy(A'), α; atol, rtol, mode)

Mooncake.TestUtils.test_rule(rng, add!, C, A; atol, rtol, mode, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, C, A, α; atol, rtol, mode, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, C, A, α, β; atol, rtol, mode)

Mooncake.TestUtils.test_rule(rng, inner, C, A; atol, rtol, mode)
Mooncake.TestUtils.test_rule(rng, inner, C', A'; atol, rtol, mode)
for α in (randn(T), One(), Zero()), β in (randn(T), One(), Zero())
Mooncake.TestUtils.test_rule(rng, scale!, C, α; atol, rtol)
Mooncake.TestUtils.test_rule(rng, scale!, C', α; atol, rtol)
Mooncake.TestUtils.test_rule(rng, scale!, C, A, α; atol, rtol)
Mooncake.TestUtils.test_rule(rng, scale!, C', A', α; atol, rtol)
Mooncake.TestUtils.test_rule(rng, scale!, copy(C'), A', α; atol, rtol)
Mooncake.TestUtils.test_rule(rng, scale!, C', copy(A'), α; atol, rtol)

Mooncake.TestUtils.test_rule(rng, add!, C, A; atol, rtol, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, C, A, α; atol, rtol, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, C, A, α, β; atol, rtol)
Comment thread
kshyatt marked this conversation as resolved.
Mooncake.TestUtils.test_rule(rng, add!, C', A'; atol, rtol, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, copy(C'), A', α; atol, rtol, is_primitive = false)
Mooncake.TestUtils.test_rule(rng, add!, C', copy(A'), α, β; atol, rtol)

Mooncake.TestUtils.test_rule(rng, inner, C, A; atol, rtol)
Mooncake.TestUtils.test_rule(rng, inner, C', A'; atol, rtol)
end
end
Loading