Small refactor

Author: mtfishman

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -1,27 +1,28 @@
 module ITensorNetworksNextTensorOperationsExt
 
 using BackendSelection: @Algorithm_str, Algorithm
-using ITensorNetworksNext: ITensorNetworksNext, contraction_order
-using ITensorNetworksNext.LazyNamedDimsArrays: symnameddims, substitute
+using ITensorNetworksNext.LazyNamedDimsArrays: LazyNamedDimsArrays, ismul, symnameddims,
+    substitute
+using ITensorNetworksNext.LazyNamedDimsArrays.TermInterface: arguments
 using NamedDimsArrays: inds
 using TensorOperations: TensorOperations, optimaltree
 
-function contraction_order_to_expr(ord)
-    return ord isa AbstractVector ? prod(contraction_order_to_expr, ord) : symnameddims(ord)
+function contraction_tree_to_expr(f, tree)
+    return if !(tree isa AbstractVector)
+        f(tree)
+    else
+        prod(Base.Fix1(contraction_tree_to_expr, f), tree)
+    end
 end
 
-function ITensorNetworksNext.contraction_order(alg::Algorithm"optimal", tn)
-    ts = [tn[i] for i in keys(tn)]
-    network = collect.(inds.(ts))
+function LazyNamedDimsArrays.optimize_contraction_order(alg::Algorithm"optimal", a)
+    @assert ismul(a)
+    ts = arguments(a)
+    inds_network = collect.(inds.(ts))
     # Converting dims to Float64 to minimize overflow issues
-    inds_to_dims = Dict(i => Float64(length(i)) for i in unique(reduce(vcat, network)))
-    order, _ = optimaltree(network, inds_to_dims)
-    # TODO: Map the integer indices back to the original tensor network vertices.
-    expr = contraction_order_to_expr(order)
-    verts = collect(keys(tn))
-    sym(i) = symnameddims(verts[i], Tuple(inds(tn[verts[i]])))
-    subs = Dict(symnameddims(i) => sym(i) for i in eachindex(verts))
-    return substitute(expr, subs)
+    inds_to_dims = Dict(i => Float64(length(i)) for i in reduce(∪, inds_network))
+    tree, _ = optimaltree(inds_network, inds_to_dims)
+    return contraction_tree_to_expr(i -> ts[i], tree)
 end
 
 end

src/LazyNamedDimsArrays/evaluation_order.jl

@@ -61,26 +61,20 @@ function flatten_expression(a)
     if !iscall(a)
         return a
     elseif ismul(a)
-        flattened_arguments = mapreduce(vcat, arguments(a)) do arg
-            return ismul(arg) ? arguments(arg) : [arg]
-        end
+        flattened_arguments = mapreduce(to_mul_arguments, vcat, arguments(a))
         return lazy(Mul(flattened_arguments))
     else
         return error("Variant not supported.")
     end
 end
 
 function optimize_evaluation_order(alg, a)
-    return optimize_evaluation_order_flattened(alg, flatten_expression(a))
-end
-
-function optimize_evaluation_order_flattened(alg, a)
     if !iscall(a)
         return a
     elseif ismul(a)
-        return optimize_contraction_order_flattened(alg, a)
+        return optimize_contraction_order(alg, a)
     else
-        # TODO: Recurse into other operations, calling `optimize_evaluation_order_flattened`.
+        # TODO: Recurse into other operations, calling `optimize_evaluation_order`.
         return error("Variant not supported.")
     end
 end
@@ -94,19 +88,21 @@ end
 using BackendSelection: @Algorithm_str, Algorithm
 default_optimize_evaluation_order_alg(a) = Algorithm"eager"()
 
-function optimize_contraction_order_flattened(alg, a)
+function optimize_contraction_order(alg, a)
     return error("`alg = $alg` not supported.")
 end
 
 using Combinatorics: combinations
-function optimize_contraction_order_flattened(alg::Algorithm"eager", a)
+function optimize_contraction_order(alg::Algorithm"eager", a)
     @assert ismul(a)
     arity(a) in (1, 2) && return a
     a1, a2 = argmin(combinations(arguments(a), 2)) do (a1, a2)
         # Penalize outer product contractions.
+        # TODO: Still order the outer products by time complexity,
+        # say by checking if there are only outer products left.
         isdisjoint(inds(a1), inds(a2)) && return typemax(Int)
         return time_complexity(*, a1, a2)
     end
     contracted_arguments = [filter(∉((a1, a2)), arguments(a)); [a1 * a2]]
-    return optimize_contraction_order_flattened(alg, lazy(Mul(contracted_arguments)))
+    return optimize_contraction_order(alg, lazy(Mul(contracted_arguments)))
 end

src/LazyNamedDimsArrays/lazyinterface.jl

@@ -157,13 +157,8 @@ 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
+    !iscall(a) ? show(io, mime, unwrap(a)) : show(io, a)
+    return nothing
 end
 add_lazy(a1, a2) = error("Not implemented.")
 sub_lazy(a) = error("Not implemented.")
@@ -179,7 +174,12 @@ function mul_lazy(a)
     end
 end
 # Note that this is nested by default.
-mul_lazy(a1, a2) = lazy(Mul([a1, a2]))
+function mul_lazy(a1, a2; flatten::Bool = false)
+    return flatten ? mul_lazy_flattened(a1, a2) : mul_lazy_nested(a1, a2)
+end
+mul_lazy_nested(a1, a2) = lazy(Mul([a1, a2]))
+to_mul_arguments(a) = ismul(a) ? arguments(a) : [a]
+mul_lazy_flattened(a1, a2) = lazy(Mul([to_mul_arguments(a1); to_mul_arguments(a2)]))
 mul_lazy(a1::Number, a2) = error("Not implemented.")
 mul_lazy(a1, a2::Number) = error("Not implemented.")
 mul_lazy(a1::Number, a2::Number) = a1 * a2

src/contract_network.jl

@@ -1,7 +1,7 @@
 using BackendSelection: @Algorithm_str, Algorithm
 using Base.Broadcast: materialize
-using ITensorNetworksNext.LazyNamedDimsArrays: lazy, optimize_evaluation_order, substitute,
-    symnameddims
+using ITensorNetworksNext.LazyNamedDimsArrays: Mul, lazy, optimize_evaluation_order,
+    substitute, symnameddims
 
 # This is related to `MatrixAlgebraKit.select_algorithm`.
 # TODO: Define this in BackendSelection.jl.
@@ -14,7 +14,9 @@ to_algorithm(alg::Algorithm; kwargs...) = merge_parameters(alg; kwargs...)
 to_algorithm(alg; kwargs...) = Algorithm(alg; kwargs...)
 
 # `contract_network`
-contract_network(alg::Algorithm, tn) = error("Not implemented.")
+function contract_network(alg::Algorithm, tn)
+    return throw(ArgumentError("`contract_network` algorithm `$(alg)` not implemented."))
+end
 function default_kwargs(::typeof(contract_network), tn)
     return (; alg = Algorithm"exact"(; order_alg = Algorithm"eager"()))
 end
@@ -23,14 +25,25 @@ function contract_network(tn; alg = default_kwargs(contract_network, tn).alg, kw
 end
 
 # `contract_network(::Algorithm"exact", ...)`
-function contract_network(alg::Algorithm"exact", tn)
-    order = @something begin
-        get(alg, :order, nothing)
-        contraction_order(
-            tn; alg = get(alg, :order_alg, default_kwargs(contraction_order, tn).alg)
-        )
+function get_order(alg::Algorithm"exact", tn)
+    # Allow specifying either `order` or `order_alg`.
+    order = get(alg, :order, nothing)
+    order = if !isnothing(order)
+        order
+    else
+        default_order_alg = default_kwargs(contraction_order, tn).alg
+        order_alg = get(alg, :order_alg, default_order_alg)
+        # TODO: Capture other keyword arguments and pass them to `contraction_order`.
+        contraction_order(tn; alg = order_alg)
     end
-    syms_to_ts = Dict(symnameddims(i, Tuple(inds(tn[i]))) => lazy(tn[i]) for i in eachindex(tn))
+    # Contraction order may or may not have indices attached, canonicalize the format
+    # by attaching indices.
+    subs = Dict(symnameddims(i) => symnameddims(i, Tuple(inds(tn[i]))) for i in keys(tn))
+    return substitute(order, subs)
+end
+function contract_network(alg::Algorithm"exact", tn)
+    order = get_order(alg, tn)
+    syms_to_ts = Dict(symnameddims(i, Tuple(inds(tn[i]))) => lazy(tn[i]) for i in keys(tn))
     tn_expression = substitute(order, syms_to_ts)
     return materialize(tn_expression)
 end
@@ -41,10 +54,16 @@ default_kwargs(::typeof(contraction_order), tn) = (; alg = Algorithm"eager"())
 function contraction_order(tn; alg = default_kwargs(contraction_order, tn).alg, kwargs...)
     return contraction_order(to_algorithm(alg; kwargs...), tn)
 end
+# Convert the tensor network to a flat symbolic multiplication expression.
+function contraction_order(alg::Algorithm"flat", tn)
+    syms = [symnameddims(i, Tuple(inds(tn[i]))) for i in keys(tn)]
+    # Same as: `reduce((a, b) -> *(a, b; flatten = true), syms)`.
+    return lazy(Mul(syms))
+end
 function contraction_order(alg::Algorithm"left_associative", tn)
     return prod(i -> symnameddims(i, Tuple(inds(tn[i]))), keys(tn))
 end
 function contraction_order(alg::Algorithm, tn)
-    s = contraction_order(tn; alg = Algorithm"left_associative"())
+    s = contraction_order(Algorithm"flat"(), tn)
     return optimize_evaluation_order(s; alg)
 end