From a7e368bfa8986dc42de4811d02a4509d1e2f8872 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 10:50:40 -0400 Subject: [PATCH 1/7] Define substitute --- Project.toml | 2 +- src/lazynameddimsarrays.jl | 76 ++++++++++++++++++++++++++++++++ test/test_lazynameddimsarrays.jl | 40 ++++++++++------- 3 files changed, 102 insertions(+), 16 deletions(-) diff --git a/Project.toml b/Project.toml index 5cfd81d3..f9a2dfd1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" authors = ["ITensor developers and contributors"] -version = "0.1.8" +version = "0.1.9" [deps] AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" diff --git a/src/lazynameddimsarrays.jl b/src/lazynameddimsarrays.jl index 4a038e86..ff3f53ac 100644 --- a/src/lazynameddimsarrays.jl +++ b/src/lazynameddimsarrays.jl @@ -39,6 +39,16 @@ function Base.show(io::IO, m::Mul) print(io, "(", join(args, " $(operation(m)) "), ")") return nothing end +function Base.hash(m::Mul, h::UInt64) + h = hash(:Mul, h) + for arg in arguments(m) + h = hash(arg, h) + end + return h +end +function map_arguments(f, m::Mul) + return Mul(map(f, arguments(m))) +end @wrapped struct LazyNamedDimsArray{ T, A <: AbstractNamedDimsArray{T}, @@ -65,6 +75,18 @@ function NamedDimsArrays.dename(a::LazyNamedDimsArray) end end +function getindex_lazy(a::AbstractArray, I...) + u = unwrap(a) + if !iscall(u) + return u[I...] + else + return error("Indexing into expression not supported.") + end +end +function Base.getindex(a::LazyNamedDimsArray, I::Int...) + return getindex_lazy(a, I...) +end + function TermInterface.arguments(a::LazyNamedDimsArray) u = unwrap(a) if !iscall(u) @@ -158,6 +180,49 @@ function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) end end +# Defined to avoid type piracy. +# TODO: Define a proper hash function +# in NamedDimsArrays.jl, maybe one that is +# independent of the order of dimensions. +function _hash(a::NamedDimsArray, h::UInt64) + h = hash(:NamedDimsArray, h) + h = hash(dename(a), h) + for i in inds(a) + h = hash(i, h) + end + return h +end +function _hash(x, h::UInt64) + return hash(x, h) +end +function Base.hash(a::LazyNamedDimsArray, h::UInt64) + h = hash(:LazyNamedDimsArray, h) + # Use `_hash`, which defines a custom hash for NamedDimsArray. + return _hash(unwrap(a), h) +end + +generic_map(f, v) = map(f, v) +generic_map(f, v::AbstractDict) = Dict(eachindex(v) .=> map(f, values(v))) +generic_map(f, v::AbstractSet) = Set([f(x) for x in v]) +function map_arguments(f, a::LazyNamedDimsArray) + u = unwrap(a) + if !iscall(u) + return error("No arguments to map.") + elseif ismul(u) + return LazyNamedDimsArray(map_arguments(f, u)) + else + return error("Variant not supported.") + end +end +function substitute(a::LazyNamedDimsArray, substitutions::AbstractDict) + haskey(substitutions, a) && return substitutions[a] + !iscall(a) && return a + return map_arguments(arg -> substitute(arg, substitutions), a) +end +function substitute(a::LazyNamedDimsArray, substitutions) + return substitute(a, Dict(substitutions)) +end + function printnode(io::IO, a::LazyNamedDimsArray) return printnode(io, unwrap(a)) end @@ -280,6 +345,17 @@ Base.size(a::SymbolicArray) = length.(axes(a)) function Base.:(==)(a::SymbolicArray, b::SymbolicArray) return symname(a) == symname(b) && axes(a) == axes(b) end +function Base.hash(a::SymbolicArray, h::UInt64) + h = hash(:SymbolicArray, h) + h = hash(symname(a), h) + return hash(size(a), h) +end +function Base.getindex(a::SymbolicArray, I...) + return error("Indexing into SymbolicArray not supported.") +end +function Base.setindex!(a::SymbolicArray, value, I...) + return error("Indexing into SymbolicArray not supported.") +end function Base.show(io::IO, mime::MIME"text/plain", a::SymbolicArray) Base.summary(io, a) println(io, ":") diff --git a/test/test_lazynameddimsarrays.jl b/test/test_lazynameddimsarrays.jl index 40735ec3..787ceb7b 100644 --- a/test/test_lazynameddimsarrays.jl +++ b/test/test_lazynameddimsarrays.jl @@ -1,8 +1,8 @@ using AbstractTrees: AbstractTrees, print_tree, printnode using Base.Broadcast: materialize using ITensorNetworksNext.LazyNamedDimsArrays: - LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, symnameddims -using NamedDimsArrays: NamedDimsArray, dename, dimnames, inds, nameddims + LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, substitute, symnameddims +using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims using TermInterface: arguments, arity, @@ -88,22 +88,32 @@ using WrappedUnions: unwrap end @testset "symnameddims" begin - a = symnameddims(:a) - b = symnameddims(:b) - c = symnameddims(:c) - @test a isa LazyNamedDimsArray - @test unwrap(a) isa NamedDimsArray - @test dename(a) isa SymbolicArray - @test dename(unwrap(a)) isa SymbolicArray - @test dename(unwrap(a)) == SymbolicArray(:a) + a1, a2, a3 = symnameddims.((:a1, :a2, :a3)) + @test a1 isa LazyNamedDimsArray + @test unwrap(a1) isa NamedDimsArray + @test dename(a1) isa SymbolicArray + @test dename(unwrap(a1)) isa SymbolicArray + @test dename(unwrap(a1)) == SymbolicArray(:a1) @test inds(a) == () - @test dimnames(a) == () + @test dimnames(a1) == () - ex = a * b * c + ex = a1 * a2 * a3 @test copy(ex) == ex - @test arguments(ex) == [a * b, c] + @test arguments(ex) == [a1 * a2, a3] @test operation(ex) ≡ * - @test sprint(show, ex) == "((a * b) * c)" - @test sprint(show, MIME"text/plain"(), ex) == "((a * b) * c)" + @test sprint(show, ex) == "((a1 * a2) * a3)" + @test sprint(show, MIME"text/plain"(), ex) == "((a1 * a2) * a3)" + end + + @testset "substitute" begin + s = symnameddims.((:a1, :a2, :a3)) + i = @names i[1:4] + a = (randn(2, 2)[i[1], i[2]], randn(2, 2)[i[2], i[3]], randn(2, 2)[i[3], i[4]]) + l = lazy.(a) + + seq = s[1] * (s[2] * s[3]) + net = substitute(seq, s .=> l) + @test net == l[1] * (l[2] * l[3]) + @test arguments(net) == [l[1], l[2] * l[3]] end end From 2716c18a16cf3c043955c8f4b4aa5bea9f5bb322 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 13:05:33 -0400 Subject: [PATCH 2/7] Code reorg --- Project.toml | 2 + src/lazynameddimsarrays.jl | 337 ++++++++++++++++++------------- test/test_lazynameddimsarrays.jl | 2 +- 3 files changed, 197 insertions(+), 144 deletions(-) diff --git a/Project.toml b/Project.toml index f9a2dfd1..b9447691 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c" +TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138" WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44" [compat] @@ -33,5 +34,6 @@ NamedGraphs = "0.6.9, 0.7" SimpleTraits = "0.9.5" SplitApplyCombine = "1.2.3" TermInterface = "2" +TypeParameterAccessors = "0.4.4" WrappedUnions = "0.3" julia = "1.10" diff --git a/src/lazynameddimsarrays.jl b/src/lazynameddimsarrays.jl index ff3f53ac..2a237b3d 100644 --- a/src/lazynameddimsarrays.jl +++ b/src/lazynameddimsarrays.jl @@ -11,70 +11,39 @@ using NamedDimsArrays: dimnames, inds using TermInterface: TermInterface, arguments, iscall, maketerm, operation, sorted_arguments +using TypeParameterAccessors: unspecify_type_parameters -# Custom version of `AbstractTrees.printnode` to -# avoid type piracy when overloading on `AbstractNamedDimsArray`. -printnode(io::IO, x) = AbstractTrees.printnode(io, x) -function printnode(io::IO, a::AbstractNamedDimsArray) - show(io, collect(dimnames(a))) - return nothing -end +lazy(x) = error("Not defined.") -struct Mul{A} - arguments::Vector{A} -end -TermInterface.arguments(m::Mul) = getfield(m, :arguments) -TermInterface.children(m::Mul) = arguments(m) -TermInterface.head(m::Mul) = operation(m) -TermInterface.iscall(m::Mul) = true -TermInterface.isexpr(m::Mul) = iscall(m) -TermInterface.maketerm(::Type{Mul}, head::typeof(*), args, metadata) = Mul(args) -TermInterface.operation(m::Mul) = * -TermInterface.sorted_arguments(m::Mul) = arguments(m) -TermInterface.sorted_children(m::Mul) = sorted_arguments(a) -ismul(x) = false -ismul(m::Mul) = true -function Base.show(io::IO, m::Mul) - args = map(arg -> sprint(printnode, arg), arguments(m)) - print(io, "(", join(args, " $(operation(m)) "), ")") - return nothing -end -function Base.hash(m::Mul, h::UInt64) - h = hash(:Mul, h) - for arg in arguments(m) - h = hash(arg, h) +generic_map(f, v) = map(f, v) +generic_map(f, v::AbstractDict) = Dict(eachindex(v) .=> map(f, values(v))) +generic_map(f, v::AbstractSet) = Set([f(x) for x in v]) + +# Defined to avoid type piracy. +# TODO: Define a proper hash function +# in NamedDimsArrays.jl, maybe one that is +# independent of the order of dimensions. +function _hash(a::NamedDimsArray, h::UInt64) + h = hash(:NamedDimsArray, h) + h = hash(dename(a), h) + for i in inds(a) + h = hash(i, h) end return h end -function map_arguments(f, m::Mul) - return Mul(map(f, arguments(m))) -end - -@wrapped struct LazyNamedDimsArray{ - T, A <: AbstractNamedDimsArray{T}, - } <: AbstractNamedDimsArray{T, Any} - union::Union{A, Mul{LazyNamedDimsArray{T, A}}} +function _hash(x, h::UInt64) + return hash(x, h) end -function NamedDimsArrays.inds(a::LazyNamedDimsArray) - u = unwrap(a) - if !iscall(u) - return inds(u) - elseif ismul(u) - return mapreduce(inds, symdiff, arguments(u)) - else - return error("Variant not supported.") - end -end -function NamedDimsArrays.dename(a::LazyNamedDimsArray) - u = unwrap(a) - if !iscall(u) - return dename(u) - else - return error("Variant not supported.") - end +# Custom version of `AbstractTrees.printnode` to +# avoid type piracy when overloading on `AbstractNamedDimsArray`. +printnode(io::IO, x) = AbstractTrees.printnode(io, x) +function printnode(io::IO, a::AbstractNamedDimsArray) + show(io, collect(dimnames(a))) + return nothing end +# Generic lazy functionality. function getindex_lazy(a::AbstractArray, I...) u = unwrap(a) if !iscall(u) @@ -83,11 +52,7 @@ function getindex_lazy(a::AbstractArray, I...) return error("Indexing into expression not supported.") end end -function Base.getindex(a::LazyNamedDimsArray, I::Int...) - return getindex_lazy(a, I...) -end - -function TermInterface.arguments(a::LazyNamedDimsArray) +function arguments_lazy(a) u = unwrap(a) if !iscall(u) return error("No arguments.") @@ -97,26 +62,19 @@ function TermInterface.arguments(a::LazyNamedDimsArray) return error("Variant not supported.") end end -function TermInterface.children(a::LazyNamedDimsArray) +function children_lazy(a) return arguments(a) end -function TermInterface.head(a::LazyNamedDimsArray) +function head_lazy(a) return operation(a) end -function TermInterface.iscall(a::LazyNamedDimsArray) +function iscall_lazy(a) return iscall(unwrap(a)) end -function TermInterface.isexpr(a::LazyNamedDimsArray) +function isexpr_lazy(a) return iscall(a) end -function TermInterface.maketerm(::Type{LazyNamedDimsArray}, head, args, metadata) - if head ≡ * - return LazyNamedDimsArray(maketerm(Mul, head, args, metadata)) - else - return error("Only mul supported right now.") - end -end -function TermInterface.operation(a::LazyNamedDimsArray) +function operation_lazy(a) u = unwrap(a) if !iscall(u) return error("No operation.") @@ -126,7 +84,7 @@ function TermInterface.operation(a::LazyNamedDimsArray) return error("Variant not supported.") end end -function TermInterface.sorted_arguments(a::LazyNamedDimsArray) +function sorted_arguments_lazy(a) u = unwrap(a) if !iscall(u) return error("No arguments.") @@ -136,28 +94,26 @@ function TermInterface.sorted_arguments(a::LazyNamedDimsArray) return error("Variant not supported.") end end -function TermInterface.sorted_children(a::LazyNamedDimsArray) +function sorted_children_lazy(a) return sorted_arguments(a) end -ismul(a::LazyNamedDimsArray) = ismul(unwrap(a)) - -function AbstractTrees.children(a::LazyNamedDimsArray) +ismul_lazy(a) = ismul(unwrap(a)) +function abstracttrees_children_lazy(a) if !iscall(a) return () else return arguments(a) end end -function AbstractTrees.nodevalue(a::LazyNamedDimsArray) +function nodevalue_lazy(a) if !iscall(a) return unwrap(a) else return operation(a) end end - using Base.Broadcast: materialize -function Base.Broadcast.materialize(a::LazyNamedDimsArray) +function materialize_lazy(a) u = unwrap(a) if !iscall(u) return u @@ -167,9 +123,8 @@ function Base.Broadcast.materialize(a::LazyNamedDimsArray) return error("Variant not supported.") end end -Base.copy(a::LazyNamedDimsArray) = materialize(a) - -function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) +copy_lazy(a) = materialize(a) +function equals_lazy(a1, a2) u1, u2 = unwrap.((a1, a2)) if !iscall(u1) && !iscall(u2) return u1 == u2 @@ -179,41 +134,174 @@ function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) return false end end +function hash_lazy(a, h::UInt64) + h = hash(Symbol(unspecify_type_parameters(typeof(a))), h) + # Use `_hash`, which defines a custom hash for NamedDimsArray. + return _hash(unwrap(a), h) +end +function map_arguments_lazy(f, a) + u = unwrap(a) + if !iscall(u) + return error("No arguments to map.") + elseif ismul(u) + return unspecify_type_parameters(typeof(a))(map_arguments(f, u)) + else + return error("Variant not supported.") + end +end -# Defined to avoid type piracy. -# TODO: Define a proper hash function -# in NamedDimsArrays.jl, maybe one that is -# independent of the order of dimensions. -function _hash(a::NamedDimsArray, h::UInt64) - h = hash(:NamedDimsArray, h) - h = hash(dename(a), h) - for i in inds(a) - h = hash(i, h) +# Lazy broadcasting. +struct LazyNamedDimsArrayStyle <: AbstractNamedDimsArrayStyle{Any} end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, f, as...) + return error("Arbitrary broadcasting not supported for LazyNamedDimsArray.") +end +# Linear operations. +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(+), a1, a2) + return a1 + a2 +end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a1, a2) + return a1 - a2 +end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), c::Number, a) + return c * a +end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a, c::Number) + return a * c +end +# Fix ambiguity error. +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a::Number, b::Number) + return a * b +end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(/), a, c::Number) + return a / c +end +function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a) + return -a +end + +# Generic functionality for Applied types, like `Mul`, `Add`, etc. +abstract type Applied end +TermInterface.head(m::Applied) = operation(m) +TermInterface.iscall(m::Applied) = true +TermInterface.isexpr(m::Applied) = iscall(m) +function Base.show(io::IO, m::Applied) + args = map(arg -> sprint(printnode, arg), arguments(m)) + print(io, "(", join(args, " $(operation(m)) "), ")") + return nothing +end +TermInterface.sorted_children(m::Applied) = sorted_arguments(m) + +struct Mul{A} <: Applied + arguments::Vector{A} +end +TermInterface.arguments(m::Mul) = getfield(m, :arguments) +TermInterface.children(m::Mul) = arguments(m) +TermInterface.maketerm(::Type{Mul}, head::typeof(*), args, metadata) = Mul(args) +TermInterface.operation(m::Mul) = * +TermInterface.sorted_arguments(m::Mul) = arguments(m) +ismul(x) = false +ismul(m::Mul) = true +function Base.hash(m::Mul, h::UInt64) + h = hash(:Mul, h) + for arg in arguments(m) + h = hash(arg, h) end return h end -function _hash(x, h::UInt64) - return hash(x, h) +function map_arguments(f, m::Mul) + return Mul(map(f, arguments(m))) end -function Base.hash(a::LazyNamedDimsArray, h::UInt64) - h = hash(:LazyNamedDimsArray, h) - # Use `_hash`, which defines a custom hash for NamedDimsArray. - return _hash(unwrap(a), h) + +@wrapped struct LazyNamedDimsArray{ + T, A <: AbstractNamedDimsArray{T}, + } <: AbstractNamedDimsArray{T, Any} + union::Union{A, Mul{LazyNamedDimsArray{T, A}}} +end +function LazyNamedDimsArray(a::AbstractNamedDimsArray) + return LazyNamedDimsArray{eltype(a), typeof(a)}(a) +end +function LazyNamedDimsArray(a::Mul{LazyNamedDimsArray{T, A}}) where {T, A} + return LazyNamedDimsArray{T, A}(a) +end +function lazy(a::AbstractNamedDimsArray) + return LazyNamedDimsArray(a) end -generic_map(f, v) = map(f, v) -generic_map(f, v::AbstractDict) = Dict(eachindex(v) .=> map(f, values(v))) -generic_map(f, v::AbstractSet) = Set([f(x) for x in v]) -function map_arguments(f, a::LazyNamedDimsArray) +function NamedDimsArrays.inds(a::LazyNamedDimsArray) u = unwrap(a) if !iscall(u) - return error("No arguments to map.") + return inds(u) elseif ismul(u) - return LazyNamedDimsArray(map_arguments(f, u)) + return mapreduce(inds, symdiff, arguments(u)) + else + return error("Variant not supported.") + end +end +function NamedDimsArrays.dename(a::LazyNamedDimsArray) + u = unwrap(a) + if !iscall(u) + return dename(u) else return error("Variant not supported.") end end +function TermInterface.maketerm(::Type{LazyNamedDimsArray}, head, args, metadata) + if head ≡ * + return LazyNamedDimsArray(maketerm(Mul, head, args, metadata)) + else + return error("Only mul supported right now.") + end +end + +# Derived functionality. +function Base.getindex(a::LazyNamedDimsArray, I::Int...) + return getindex_lazy(a, I...) +end +function TermInterface.arguments(a::LazyNamedDimsArray) + return arguments_lazy(a) +end +function TermInterface.children(a::LazyNamedDimsArray) + return children_lazy(a) +end +function TermInterface.head(a::LazyNamedDimsArray) + return head_lazy(a) +end +function TermInterface.iscall(a::LazyNamedDimsArray) + return iscall_lazy(a) +end +function TermInterface.isexpr(a::LazyNamedDimsArray) + return isexpr_lazy(a) +end +function TermInterface.operation(a::LazyNamedDimsArray) + return operation_lazy(a) +end +function TermInterface.sorted_arguments(a::LazyNamedDimsArray) + return sorted_arguments_lazy(a) +end +function AbstractTrees.children(a::LazyNamedDimsArray) + return abstracttrees_children_lazy(a) +end +function TermInterface.sorted_children(a::LazyNamedDimsArray) + return sorted_children_lazy(a) +end +ismul(a::LazyNamedDimsArray) = ismul_lazy(a) +function AbstractTrees.nodevalue(a::LazyNamedDimsArray) + return nodevalue_lazy(a) +end +function Base.Broadcast.materialize(a::LazyNamedDimsArray) + return materialize_lazy(a) +end +Base.copy(a::LazyNamedDimsArray) = copy_lazy(a) +function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) + return equals_lazy(a1, a2) +end +function Base.hash(a::LazyNamedDimsArray, h::UInt64) + return hash_lazy(a, h) +end +function map_arguments(f, a::LazyNamedDimsArray) + return map_arguments_lazy(f, a) +end + function substitute(a::LazyNamedDimsArray, substitutions::AbstractDict) haskey(substitutions, a) && return substitutions[a] !iscall(a) && return a @@ -280,47 +368,10 @@ function Base.:-(a::LazyNamedDimsArray) return error("Not implemented.") end -function LazyNamedDimsArray(a::AbstractNamedDimsArray) - return LazyNamedDimsArray{eltype(a), typeof(a)}(a) -end -function LazyNamedDimsArray(a::Mul{LazyNamedDimsArray{T, A}}) where {T, A} - return LazyNamedDimsArray{T, A}(a) -end -function lazy(a::AbstractNamedDimsArray) - return LazyNamedDimsArray(a) -end - # Broadcasting -struct LazyNamedDimsArrayStyle <: AbstractNamedDimsArrayStyle{Any} end function Base.BroadcastStyle(::Type{<:LazyNamedDimsArray}) return LazyNamedDimsArrayStyle() end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, f, as...) - return error("Arbitrary broadcasting not supported for LazyNamedDimsArray.") -end -# Linear operations. -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(+), a1, a2) - return a1 + a2 -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a1, a2) - return a1 - a2 -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), c::Number, a) - return c * a -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a, c::Number) - return a * c -end -# Fix ambiguity error. -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a::Number, b::Number) - return a * b -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(/), a, c::Number) - return a / c -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a) - return -a -end struct SymbolicArray{T, N, Name, Axes <: NTuple{N, AbstractUnitRange{<:Integer}}} <: AbstractArray{T, N} name::Name @@ -350,10 +401,10 @@ function Base.hash(a::SymbolicArray, h::UInt64) h = hash(symname(a), h) return hash(size(a), h) end -function Base.getindex(a::SymbolicArray, I...) +function Base.getindex(a::SymbolicArray{<:Any, N}, I::Vararg{Int, N}) where {N} return error("Indexing into SymbolicArray not supported.") end -function Base.setindex!(a::SymbolicArray, value, I...) +function Base.setindex!(a::SymbolicArray{<:Any, N}, value, I::Vararg{Int, N}) where {N} return error("Indexing into SymbolicArray not supported.") end function Base.show(io::IO, mime::MIME"text/plain", a::SymbolicArray) diff --git a/test/test_lazynameddimsarrays.jl b/test/test_lazynameddimsarrays.jl index 787ceb7b..996154e1 100644 --- a/test/test_lazynameddimsarrays.jl +++ b/test/test_lazynameddimsarrays.jl @@ -94,7 +94,7 @@ using WrappedUnions: unwrap @test dename(a1) isa SymbolicArray @test dename(unwrap(a1)) isa SymbolicArray @test dename(unwrap(a1)) == SymbolicArray(:a1) - @test inds(a) == () + @test inds(a1) == () @test dimnames(a1) == () ex = a1 * a2 * a3 From 96a9ebf2c00eb7d3ee62948a0c11266120ee708b Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 13:50:02 -0400 Subject: [PATCH 3/7] More reorg --- src/lazynameddimsarrays.jl | 73 ++++++++++++++++++++------------ test/test_lazynameddimsarrays.jl | 15 ++++--- 2 files changed, 54 insertions(+), 34 deletions(-) diff --git a/src/lazynameddimsarrays.jl b/src/lazynameddimsarrays.jl index 2a237b3d..bc977c76 100644 --- a/src/lazynameddimsarrays.jl +++ b/src/lazynameddimsarrays.jl @@ -37,8 +37,8 @@ end # Custom version of `AbstractTrees.printnode` to # avoid type piracy when overloading on `AbstractNamedDimsArray`. -printnode(io::IO, x) = AbstractTrees.printnode(io, x) -function printnode(io::IO, a::AbstractNamedDimsArray) +printnode_nameddims(io::IO, x) = AbstractTrees.printnode(io, x) +function printnode_nameddims(io::IO, a::AbstractNamedDimsArray) show(io, collect(dimnames(a))) return nothing end @@ -149,6 +149,37 @@ function map_arguments_lazy(f, a) return error("Variant not supported.") end end +function substitute_lazy(a, substitutions::AbstractDict) + haskey(substitutions, a) && return substitutions[a] + !iscall(a) && return a + return map_arguments(arg -> substitute(arg, substitutions), a) +end +function substitute_lazy(a, substitutions) + return substitute(a, Dict(substitutions)) +end +function printnode_lazy(io, a) + # Use `printnode_nameddims` to avoid type piracy, + # since it overloads on `AbstractNamedDimsArray`. + return printnode_nameddims(io, unwrap(a)) +end +function show_lazy(io::IO, a) + if !iscall(a) + return show(io, unwrap(a)) + else + return AbstractTrees.printnode(io, a) + end +end +function show_lazy(io::IO, mime::MIME"text/plain", a) + summary(io, a) + println(io, ":") + if !iscall(a) + show(io, mime, unwrap(a)) + return nothing + else + show(io, a) + return nothing + end +end # Lazy broadcasting. struct LazyNamedDimsArrayStyle <: AbstractNamedDimsArrayStyle{Any} end @@ -185,7 +216,7 @@ TermInterface.head(m::Applied) = operation(m) TermInterface.iscall(m::Applied) = true TermInterface.isexpr(m::Applied) = iscall(m) function Base.show(io::IO, m::Applied) - args = map(arg -> sprint(printnode, arg), arguments(m)) + args = map(arg -> sprint(AbstractTrees.printnode, arg), arguments(m)) print(io, "(", join(args, " $(operation(m)) "), ")") return nothing end @@ -301,37 +332,20 @@ end function map_arguments(f, a::LazyNamedDimsArray) return map_arguments_lazy(f, a) end - -function substitute(a::LazyNamedDimsArray, substitutions::AbstractDict) - haskey(substitutions, a) && return substitutions[a] - !iscall(a) && return a - return map_arguments(arg -> substitute(arg, substitutions), a) -end function substitute(a::LazyNamedDimsArray, substitutions) - return substitute(a, Dict(substitutions)) -end - -function printnode(io::IO, a::LazyNamedDimsArray) - return printnode(io, unwrap(a)) + return substitute_lazy(a, substitutions) end function AbstractTrees.printnode(io::IO, a::LazyNamedDimsArray) - return printnode(io, a) + return printnode_lazy(io, a) +end +function printnode_nameddims(io::IO, a::LazyNamedDimsArray) + return printnode_lazy(io, a) end function Base.show(io::IO, a::LazyNamedDimsArray) - if !iscall(a) - return show(io, unwrap(a)) - else - return printnode(io, a) - end + return show_lazy(io, a) end function Base.show(io::IO, mime::MIME"text/plain", a::LazyNamedDimsArray) - if !iscall(a) - @invoke show(io, mime, a::AbstractNamedDimsArray) - return nothing - else - show(io, a) - return nothing - end + return show_lazy(io, mime, a) end function Base.:*(a::LazyNamedDimsArray) @@ -427,13 +441,16 @@ const SymbolicNamedDimsArray{T, N, Parent <: SymbolicArray{T, N}, DimNames} = function symnameddims(name) return lazy(NamedDimsArray(SymbolicArray(name), ())) end -function printnode(io::IO, a::SymbolicNamedDimsArray) +function AbstractTrees.printnode(io::IO, a::SymbolicNamedDimsArray) print(io, symname(dename(a))) if ndims(a) > 0 print(io, "[", join(dimnames(a), ","), "]") end return nothing end +function printnode_nameddims(io::IO, a::SymbolicNamedDimsArray) + return AbstractTrees.printnode(io, a) +end function Base.:(==)(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray) return issetequal(inds(a), inds(b)) && dename(a) == dename(b) end diff --git a/test/test_lazynameddimsarrays.jl b/test/test_lazynameddimsarrays.jl index 996154e1..a9ec559f 100644 --- a/test/test_lazynameddimsarrays.jl +++ b/test/test_lazynameddimsarrays.jl @@ -2,7 +2,7 @@ using AbstractTrees: AbstractTrees, print_tree, printnode using Base.Broadcast: materialize using ITensorNetworksNext.LazyNamedDimsArrays: LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, substitute, symnameddims -using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims +using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims, namedoneto using TermInterface: arguments, arity, @@ -19,9 +19,10 @@ using WrappedUnions: unwrap @testset "LazyNamedDimsArrays" begin @testset "Basics" begin - a1 = nameddims(randn(2, 2), (:i, :j)) - a2 = nameddims(randn(2, 2), (:j, :k)) - a3 = nameddims(randn(2, 2), (:k, :l)) + i, j, k, l = namedoneto.(2, (:i, :j, :k, :l)) + a1 = randn(i, j) + a2 = randn(j, k) + a3 = randn(k, l) l1, l2, l3 = lazy.((a1, a2, a3)) for li in (l1, l2, l3) @test li isa LazyNamedDimsArray @@ -81,7 +82,8 @@ using WrappedUnions: unwrap @test AbstractTrees.children(l) == [l1 * l2, l3] @test AbstractTrees.nodevalue(l) ≡ * @test sprint(show, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" - @test sprint(show, MIME"text/plain"(), l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" + @test sprint(show, MIME"text/plain"(), l) == + "named(Base.OneTo(2), :i)×named(Base.OneTo(2), :l) LazyNamedDimsArray{Float64, …}:\n(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(printnode, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(print_tree, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])\n├─ ([:i, :j] * [:j, :k])\n│ ├─ [:i, :j]\n│ └─ [:j, :k]\n└─ [:k, :l]\n" @@ -102,7 +104,8 @@ using WrappedUnions: unwrap @test arguments(ex) == [a1 * a2, a3] @test operation(ex) ≡ * @test sprint(show, ex) == "((a1 * a2) * a3)" - @test sprint(show, MIME"text/plain"(), ex) == "((a1 * a2) * a3)" + @test sprint(show, MIME"text/plain"(), ex) == + "0-dimensional LazyNamedDimsArray{Any, …}:\n((a1 * a2) * a3)" end @testset "substitute" begin From 6ddfdf1063c7a2a9c90f7b4c7d9494c7bebcfe67 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 15:01:14 -0400 Subject: [PATCH 4/7] More reorg --- src/lazynameddimsarrays.jl | 316 ++++++++++++++++--------------------- 1 file changed, 133 insertions(+), 183 deletions(-) diff --git a/src/lazynameddimsarrays.jl b/src/lazynameddimsarrays.jl index bc977c76..f7cf86da 100644 --- a/src/lazynameddimsarrays.jl +++ b/src/lazynameddimsarrays.jl @@ -44,6 +44,13 @@ function printnode_nameddims(io::IO, a::AbstractNamedDimsArray) end # Generic lazy functionality. +function maketerm_lazy(type::Type, head, args, metadata) + if head ≡ * + return type(maketerm(Mul, head, args, metadata)) + else + return error("Only mul supported right now.") + end +end function getindex_lazy(a::AbstractArray, I...) u = unwrap(a) if !iscall(u) @@ -144,7 +151,7 @@ function map_arguments_lazy(f, a) if !iscall(u) return error("No arguments to map.") elseif ismul(u) - return unspecify_type_parameters(typeof(a))(map_arguments(f, u)) + return lazy(map_arguments(f, u)) else return error("Variant not supported.") end @@ -180,6 +187,44 @@ function show_lazy(io::IO, mime::MIME"text/plain", a) return nothing end end +add_lazy(a1, a2) = error("Not implemented.") +sub_lazy(a) = error("Not implemented.") +sub_lazy(a1, a2) = error("Not implemented.") +function mul_lazy(a) + u = unwrap(a) + if !iscall(u) + return lazy(Mul([a])) + elseif ismul(u) + return a + else + return error("Variant not supported.") + end +end +# Note that this is nested by default. +mul_lazy(a1, a2) = lazy(Mul([a1, a2])) +mul_lazy(c::Number, a) = error("Not implemented.") +mul_lazy(a, c::Number) = error("Not implemented.") +div_lazy(a, c::Number) = error("Not implemented.") + +# NamedDimsArrays.jl interface. +function inds_lazy(a) + u = unwrap(a) + if !iscall(u) + return inds(u) + elseif ismul(u) + return mapreduce(inds, symdiff, arguments(u)) + else + return error("Variant not supported.") + end +end +function dename_lazy(a) + u = unwrap(a) + if !iscall(u) + return dename(u) + else + return error("Variant not supported.") + end +end # Lazy broadcasting. struct LazyNamedDimsArrayStyle <: AbstractNamedDimsArrayStyle{Any} end @@ -187,61 +232,60 @@ function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, f, as...) return error("Arbitrary broadcasting not supported for LazyNamedDimsArray.") end # Linear operations. -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(+), a1, a2) - return a1 + a2 -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a1, a2) - return a1 - a2 -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), c::Number, a) - return c * a -end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a, c::Number) - return a * c -end -# Fix ambiguity error. -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a::Number, b::Number) - return a * b +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(+), a1, a2) = a1 + a2 +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a1, a2) = a1 - a2 +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), c::Number, a) = c * a +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a, c::Number) = a * c +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(*), a::Number, b::Number) = a * b +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(/), a, c::Number) = a / c +Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a) = -a + +# Generic functionality for Applied types, like `Mul`, `Add`, etc. +ismul(a) = operation(a) ≡ * +head_applied(a) = operation(a) +iscall_applied(a) = true +isexpr_applied(a) = iscall(a) +function show_applied(io::IO, a) + args = map(arg -> sprint(AbstractTrees.printnode, arg), arguments(a)) + print(io, "(", join(args, " $(operation(a)) "), ")") + return nothing end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(/), a, c::Number) - return a / c +sorted_arguments_applied(a) = arguments(a) +children_applied(a) = arguments(a) +sorted_children_applied(a) = sorted_arguments(a) +function maketerm_applied(type, head, args, metadata) + term = type(args) + @assert head ≡ operation(term) + return term end -function Broadcast.broadcasted(::LazyNamedDimsArrayStyle, ::typeof(-), a) - return -a +map_arguments_applied(f, a) = unspecify_type_parameters(typeof(a))(map(f, arguments(a))) +function hash_applied(a, h::UInt64) + h = hash(Symbol(unspecify_type_parameters(typeof(a))), h) + for arg in arguments(a) + h = hash(arg, h) + end + return h end -# Generic functionality for Applied types, like `Mul`, `Add`, etc. abstract type Applied end -TermInterface.head(m::Applied) = operation(m) -TermInterface.iscall(m::Applied) = true -TermInterface.isexpr(m::Applied) = iscall(m) -function Base.show(io::IO, m::Applied) - args = map(arg -> sprint(AbstractTrees.printnode, arg), arguments(m)) - print(io, "(", join(args, " $(operation(m)) "), ")") - return nothing -end -TermInterface.sorted_children(m::Applied) = sorted_arguments(m) +TermInterface.head(a::Applied) = head_applied(a) +TermInterface.iscall(a::Applied) = iscall_applied(a) +TermInterface.isexpr(a::Applied) = isexpr_applied(a) +Base.show(io::IO, a::Applied) = show_applied(io, a) +TermInterface.sorted_arguments(a::Applied) = sorted_arguments_applied(a) +TermInterface.children(a::Applied) = children_applied(a) +TermInterface.sorted_children(a::Applied) = sorted_children_applied(a) +function TermInterface.maketerm(type::Type{<:Applied}, head, args, metadata) + return maketerm_applied(type, head, args, metadata) +end +map_arguments(f, a::Applied) = map_arguments_applied(f, a) +Base.hash(a::Applied, h::UInt64) = hash_applied(a, h) struct Mul{A} <: Applied arguments::Vector{A} end TermInterface.arguments(m::Mul) = getfield(m, :arguments) -TermInterface.children(m::Mul) = arguments(m) -TermInterface.maketerm(::Type{Mul}, head::typeof(*), args, metadata) = Mul(args) TermInterface.operation(m::Mul) = * -TermInterface.sorted_arguments(m::Mul) = arguments(m) -ismul(x) = false -ismul(m::Mul) = true -function Base.hash(m::Mul, h::UInt64) - h = hash(:Mul, h) - for arg in arguments(m) - h = hash(arg, h) - end - return h -end -function map_arguments(f, m::Mul) - return Mul(map(f, arguments(m))) -end @wrapped struct LazyNamedDimsArray{ T, A <: AbstractNamedDimsArray{T}, @@ -254,139 +298,53 @@ end function LazyNamedDimsArray(a::Mul{LazyNamedDimsArray{T, A}}) where {T, A} return LazyNamedDimsArray{T, A}(a) end -function lazy(a::AbstractNamedDimsArray) - return LazyNamedDimsArray(a) -end - -function NamedDimsArrays.inds(a::LazyNamedDimsArray) - u = unwrap(a) - if !iscall(u) - return inds(u) - elseif ismul(u) - return mapreduce(inds, symdiff, arguments(u)) - else - return error("Variant not supported.") - end -end -function NamedDimsArrays.dename(a::LazyNamedDimsArray) - u = unwrap(a) - if !iscall(u) - return dename(u) - else - return error("Variant not supported.") - end -end -function TermInterface.maketerm(::Type{LazyNamedDimsArray}, head, args, metadata) - if head ≡ * - return LazyNamedDimsArray(maketerm(Mul, head, args, metadata)) - else - return error("Only mul supported right now.") - end -end - -# Derived functionality. -function Base.getindex(a::LazyNamedDimsArray, I::Int...) - return getindex_lazy(a, I...) -end -function TermInterface.arguments(a::LazyNamedDimsArray) - return arguments_lazy(a) -end -function TermInterface.children(a::LazyNamedDimsArray) - return children_lazy(a) -end -function TermInterface.head(a::LazyNamedDimsArray) - return head_lazy(a) -end -function TermInterface.iscall(a::LazyNamedDimsArray) - return iscall_lazy(a) -end -function TermInterface.isexpr(a::LazyNamedDimsArray) - return isexpr_lazy(a) -end -function TermInterface.operation(a::LazyNamedDimsArray) - return operation_lazy(a) -end -function TermInterface.sorted_arguments(a::LazyNamedDimsArray) - return sorted_arguments_lazy(a) -end -function AbstractTrees.children(a::LazyNamedDimsArray) - return abstracttrees_children_lazy(a) -end -function TermInterface.sorted_children(a::LazyNamedDimsArray) - return sorted_children_lazy(a) -end -ismul(a::LazyNamedDimsArray) = ismul_lazy(a) -function AbstractTrees.nodevalue(a::LazyNamedDimsArray) - return nodevalue_lazy(a) -end -function Base.Broadcast.materialize(a::LazyNamedDimsArray) - return materialize_lazy(a) -end -Base.copy(a::LazyNamedDimsArray) = copy_lazy(a) -function Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) - return equals_lazy(a1, a2) -end -function Base.hash(a::LazyNamedDimsArray, h::UInt64) - return hash_lazy(a, h) -end -function map_arguments(f, a::LazyNamedDimsArray) - return map_arguments_lazy(f, a) -end -function substitute(a::LazyNamedDimsArray, substitutions) - return substitute_lazy(a, substitutions) -end -function AbstractTrees.printnode(io::IO, a::LazyNamedDimsArray) - return printnode_lazy(io, a) -end -function printnode_nameddims(io::IO, a::LazyNamedDimsArray) - return printnode_lazy(io, a) -end -function Base.show(io::IO, a::LazyNamedDimsArray) - return show_lazy(io, a) -end -function Base.show(io::IO, mime::MIME"text/plain", a::LazyNamedDimsArray) - return show_lazy(io, mime, a) -end - -function Base.:*(a::LazyNamedDimsArray) - u = unwrap(a) - if !iscall(u) - return LazyNamedDimsArray(Mul([lazy(u)])) - elseif ismul(u) - return a - else - return error("Variant not supported.") - end -end +lazy(a::LazyNamedDimsArray) = a +lazy(a::AbstractNamedDimsArray) = LazyNamedDimsArray(a) +lazy(a::Mul{<:LazyNamedDimsArray}) = LazyNamedDimsArray(a) -function Base.:*(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) - # Nested by default. - return LazyNamedDimsArray(Mul([a1, a2])) -end -function Base.:+(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) - return error("Not implemented.") -end -function Base.:-(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) - return error("Not implemented.") -end -function Base.:*(c::Number, a::LazyNamedDimsArray) - return error("Not implemented.") -end -function Base.:*(a::LazyNamedDimsArray, c::Number) - return error("Not implemented.") -end -function Base.:/(a::LazyNamedDimsArray, c::Number) - return error("Not implemented.") -end -function Base.:-(a::LazyNamedDimsArray) - return error("Not implemented.") -end +NamedDimsArrays.inds(a::LazyNamedDimsArray) = inds_lazy(a) +NamedDimsArrays.dename(a::LazyNamedDimsArray) = dename_lazy(a) # Broadcasting function Base.BroadcastStyle(::Type{<:LazyNamedDimsArray}) return LazyNamedDimsArrayStyle() end +# Derived functionality. +function TermInterface.maketerm(type::Type{LazyNamedDimsArray}, head, args, metadata) + return maketerm_lazy(type, head, args, metadata) +end +Base.getindex(a::LazyNamedDimsArray, I::Int...) = getindex_lazy(a, I...) +TermInterface.arguments(a::LazyNamedDimsArray) = arguments_lazy(a) +TermInterface.children(a::LazyNamedDimsArray) = children_lazy(a) +TermInterface.head(a::LazyNamedDimsArray) = head_lazy(a) +TermInterface.iscall(a::LazyNamedDimsArray) = iscall_lazy(a) +TermInterface.isexpr(a::LazyNamedDimsArray) = isexpr_lazy(a) +TermInterface.operation(a::LazyNamedDimsArray) = operation_lazy(a) +TermInterface.sorted_arguments(a::LazyNamedDimsArray) = sorted_arguments_lazy(a) +AbstractTrees.children(a::LazyNamedDimsArray) = abstracttrees_children_lazy(a) +TermInterface.sorted_children(a::LazyNamedDimsArray) = sorted_children_lazy(a) +ismul(a::LazyNamedDimsArray) = ismul_lazy(a) +AbstractTrees.nodevalue(a::LazyNamedDimsArray) = nodevalue_lazy(a) +Base.Broadcast.materialize(a::LazyNamedDimsArray) = materialize_lazy(a) +Base.copy(a::LazyNamedDimsArray) = copy_lazy(a) +Base.:(==)(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = equals_lazy(a1, a2) +Base.hash(a::LazyNamedDimsArray, h::UInt64) = hash_lazy(a, h) +map_arguments(f, a::LazyNamedDimsArray) = map_arguments_lazy(f, a) +substitute(a::LazyNamedDimsArray, substitutions) = substitute_lazy(a, substitutions) +AbstractTrees.printnode(io::IO, a::LazyNamedDimsArray) = printnode_lazy(io, a) +printnode_nameddims(io::IO, a::LazyNamedDimsArray) = printnode_lazy(io, a) +Base.show(io::IO, a::LazyNamedDimsArray) = show_lazy(io, a) +Base.show(io::IO, mime::MIME"text/plain", a::LazyNamedDimsArray) = show_lazy(io, mime, a) +Base.:*(a::LazyNamedDimsArray) = mul_lazy(a) +Base.:*(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = mul_lazy(a1, a2) +Base.:+(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = add_lazy(a1, a2) +Base.:-(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = sub_lazy(a1, a2) +Base.:*(c::Number, a::LazyNamedDimsArray) = mul_lazy(c, a) +Base.:*(a::LazyNamedDimsArray, c::Number) = mul_lazy(a, c) +Base.:/(a::LazyNamedDimsArray, c::Number) = div_lazy(a, c) +Base.:-(a::LazyNamedDimsArray) = sub_lazy(a) + struct SymbolicArray{T, N, Name, Axes <: NTuple{N, AbstractUnitRange{<:Integer}}} <: AbstractArray{T, N} name::Name axes::Axes @@ -448,20 +406,12 @@ function AbstractTrees.printnode(io::IO, a::SymbolicNamedDimsArray) end return nothing end -function printnode_nameddims(io::IO, a::SymbolicNamedDimsArray) - return AbstractTrees.printnode(io, a) -end +printnode_nameddims(io::IO, a::SymbolicNamedDimsArray) = AbstractTrees.printnode(io, a) function Base.:(==)(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray) return issetequal(inds(a), inds(b)) && dename(a) == dename(b) end -function Base.:*(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray) - return lazy(a) * lazy(b) -end -function Base.:*(a::SymbolicNamedDimsArray, b::LazyNamedDimsArray) - return lazy(a) * b -end -function Base.:*(a::LazyNamedDimsArray, b::SymbolicNamedDimsArray) - return a * lazy(b) -end +Base.:*(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray) = lazy(a) * lazy(b) +Base.:*(a::SymbolicNamedDimsArray, b::LazyNamedDimsArray) = lazy(a) * b +Base.:*(a::LazyNamedDimsArray, b::SymbolicNamedDimsArray) = a * lazy(b) end From 10a0e0054e460d12ed0697552109984cf1048ccc Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 18:11:01 -0400 Subject: [PATCH 5/7] Fix ambiguity --- src/lazynameddimsarrays.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/lazynameddimsarrays.jl b/src/lazynameddimsarrays.jl index f7cf86da..e1b4b270 100644 --- a/src/lazynameddimsarrays.jl +++ b/src/lazynameddimsarrays.jl @@ -202,9 +202,10 @@ function mul_lazy(a) end # Note that this is nested by default. mul_lazy(a1, a2) = lazy(Mul([a1, a2])) -mul_lazy(c::Number, a) = error("Not implemented.") -mul_lazy(a, c::Number) = error("Not implemented.") -div_lazy(a, c::Number) = error("Not implemented.") +mul_lazy(a1::Number, a2) = error("Not implemented.") +mul_lazy(a1, a2::Number) = error("Not implemented.") +mul_lazy(a1::Number, a2::Number) = a1 * a2 +div_lazy(a1, a2::Number) = error("Not implemented.") # NamedDimsArrays.jl interface. function inds_lazy(a) @@ -340,9 +341,9 @@ Base.:*(a::LazyNamedDimsArray) = mul_lazy(a) Base.:*(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = mul_lazy(a1, a2) Base.:+(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = add_lazy(a1, a2) Base.:-(a1::LazyNamedDimsArray, a2::LazyNamedDimsArray) = sub_lazy(a1, a2) -Base.:*(c::Number, a::LazyNamedDimsArray) = mul_lazy(c, a) -Base.:*(a::LazyNamedDimsArray, c::Number) = mul_lazy(a, c) -Base.:/(a::LazyNamedDimsArray, c::Number) = div_lazy(a, c) +Base.:*(a1::Number, a2::LazyNamedDimsArray) = mul_lazy(a1, a2) +Base.:*(a1::LazyNamedDimsArray, a2::Number) = mul_lazy(a1, a2) +Base.:/(a1::LazyNamedDimsArray, a2::Number) = div_lazy(a1, a2) Base.:-(a::LazyNamedDimsArray) = sub_lazy(a) struct SymbolicArray{T, N, Name, Axes <: NTuple{N, AbstractUnitRange{<:Integer}}} <: AbstractArray{T, N} From 4ebbfa548140fffda974fd7932e96e2d80fbecb7 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 19:27:27 -0400 Subject: [PATCH 6/7] Fix tests --- test/test_lazynameddimsarrays.jl | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/test/test_lazynameddimsarrays.jl b/test/test_lazynameddimsarrays.jl index a9ec559f..463ff68a 100644 --- a/test/test_lazynameddimsarrays.jl +++ b/test/test_lazynameddimsarrays.jl @@ -1,23 +1,18 @@ using AbstractTrees: AbstractTrees, print_tree, printnode using Base.Broadcast: materialize -using ITensorNetworksNext.LazyNamedDimsArrays: - LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, substitute, symnameddims +using ITensorNetworksNext.LazyNamedDimsArrays: LazyNamedDimsArrays, LazyNamedDimsArray, + Mul, SymbolicArray, ismul, lazy, substitute, symnameddims using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims, namedoneto -using TermInterface: - arguments, - arity, - children, - head, - iscall, - isexpr, - maketerm, - operation, - sorted_arguments, - sorted_children +using TermInterface: arguments, arity, children, head, iscall, isexpr, maketerm, operation, + sorted_arguments, sorted_children using Test: @test, @test_throws, @testset using WrappedUnions: unwrap @testset "LazyNamedDimsArrays" begin + function sprint_namespaced(x) + context = (:module => LazyNamedDimsArrays) + return sprint(show, MIME"text/plain"(), x; context) + end @testset "Basics" begin i, j, k, l = namedoneto.(2, (:i, :j, :k, :l)) a1 = randn(i, j) @@ -63,8 +58,8 @@ using WrappedUnions: unwrap @test sprint(show, l1) == sprint(show, a1) # TODO: Fix this test, it is basically correct but the type parameters # print in a different way. - # @test sprint(show, MIME"text/plain"(), l1) == - # replace(sprint(show, MIME"text/plain"(), a1), "NamedDimsArray" => "LazyNamedDimsArray") + # @test sprint_namespaced(l1) == + # replace(sprint_namespaced(a1), "NamedDimsArray" => "LazyNamedDimsArray") @test sprint(printnode, l1) == "[:i, :j]" @test sprint(print_tree, l1) == "[:i, :j]\n" @@ -82,7 +77,7 @@ using WrappedUnions: unwrap @test AbstractTrees.children(l) == [l1 * l2, l3] @test AbstractTrees.nodevalue(l) ≡ * @test sprint(show, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" - @test sprint(show, MIME"text/plain"(), l) == + @test sprint_namespaced(l) == "named(Base.OneTo(2), :i)×named(Base.OneTo(2), :l) LazyNamedDimsArray{Float64, …}:\n(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(printnode, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(print_tree, l) == @@ -104,7 +99,7 @@ using WrappedUnions: unwrap @test arguments(ex) == [a1 * a2, a3] @test operation(ex) ≡ * @test sprint(show, ex) == "((a1 * a2) * a3)" - @test sprint(show, MIME"text/plain"(), ex) == + @test sprint_namespaced(ex) == "0-dimensional LazyNamedDimsArray{Any, …}:\n((a1 * a2) * a3)" end From c0079563bee8a9feca0acb4c06238d88c8b22ffb Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 8 Oct 2025 19:36:03 -0400 Subject: [PATCH 7/7] Fix show tests --- test/test_lazynameddimsarrays.jl | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/test/test_lazynameddimsarrays.jl b/test/test_lazynameddimsarrays.jl index 463ff68a..cc86fdce 100644 --- a/test/test_lazynameddimsarrays.jl +++ b/test/test_lazynameddimsarrays.jl @@ -11,7 +11,8 @@ using WrappedUnions: unwrap @testset "LazyNamedDimsArrays" begin function sprint_namespaced(x) context = (:module => LazyNamedDimsArrays) - return sprint(show, MIME"text/plain"(), x; context) + module_prefix = "ITensorNetworksNext.LazyNamedDimsArrays." + return replace(sprint(show, MIME"text/plain"(), x; context), module_prefix => "") end @testset "Basics" begin i, j, k, l = namedoneto.(2, (:i, :j, :k, :l)) @@ -78,10 +79,14 @@ using WrappedUnions: unwrap @test AbstractTrees.nodevalue(l) ≡ * @test sprint(show, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint_namespaced(l) == - "named(Base.OneTo(2), :i)×named(Base.OneTo(2), :l) LazyNamedDimsArray{Float64, …}:\n(([:i, :j] * [:j, :k]) * [:k, :l])" + "named(Base.OneTo(2), :i)×named(Base.OneTo(2), :l) " * + "LazyNamedDimsArray{Float64, …}:\n(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(printnode, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])" @test sprint(print_tree, l) == - "(([:i, :j] * [:j, :k]) * [:k, :l])\n├─ ([:i, :j] * [:j, :k])\n│ ├─ [:i, :j]\n│ └─ [:j, :k]\n└─ [:k, :l]\n" + "(([:i, :j] * [:j, :k]) * [:k, :l])\n" * + "├─ ([:i, :j] * [:j, :k])\n" * + "│ ├─ [:i, :j]\n│ └─ [:j, :k]\n" * + "└─ [:k, :l]\n" end @testset "symnameddims" begin