handle cases with more tensors

Author: lkdvos

src/indexnotation/optimaltree.jl

@@ -4,10 +4,11 @@ function optimaltree(network, optdata::Dict; verbose::Bool=false)
     numindices = length(allindices)
     costtype = valtype(optdata)
     allcosts = costtype[get(optdata, i, one(costtype)) for i in allindices]
-    maxcost = addcost(mulcost(reduce(mulcost, allcosts; init=one(costtype)),
-                              maximum(allcosts)), one(costtype))
-    # add one for type stability: Power -> Poly
-    # and for dealing with cases where all sizes are 1 -> maxcost = 2
+    maxcost = max(mulcost(reduce(mulcost, allcosts; init=one(costtype)),
+                          maximum(allcosts)), costtype(length(network) - 1))
+    # add zero for type stability: Power -> Poly
+    maxcost = addcost(maxcost, zero(costtype))
+
     tensorcosts = Vector{costtype}(undef, numtensors)
     for k in 1:numtensors
         tensorcosts[k] = mapreduce(i -> get(optdata, i, one(costtype)), mulcost, network[k];

test/tensoropt.jl

@@ -58,8 +58,13 @@ end
 
 @testset "Issue #206" begin
     # https://github.com/Jutho/TensorOperations.jl/issues/206
-    network = [[:a, :b], [:a], [:b]]
-    opt_data = Dict(:a => 1, :b => 1)
+    network = [[Symbol(1), Symbol(2)], [Symbol(1)], [Symbol(2)]]
+    opt_data = Dict(Symbol(1) => 1, Symbol(2) => 1)
     tree, cost = TensorOperations.optimaltree(network, opt_data)
     @test cost == 2
+
+    network = network = [[:a, :b, :c], [:a], [:b], [:c]]
+    opt_data = Dict(:a => 1, :b => 1, :c => 1)
+    tree, cost = TensorOperations.optimaltree(network, opt_data)
+    @test cost == 3
 end