update

Author: GiggleLiu

Project.toml

@@ -28,7 +28,7 @@ Compose = "0.9"
 FFTW = "1.4"
 LightGraphs = "1.3"
 Mods = "1.3"
-OMEinsum = "0.4"
+OMEinsum = "0.4, 0.5"
 OMEinsumContractionOrders = "0.2, 0.3, 0.4"
 Polynomials = "2.0"
 Primes = "0.5"

src/GraphTensorNetworks.jl

@@ -11,6 +11,38 @@ export timespace_complexity, @ein_str
 
 project_relative_path(xs...) = normpath(joinpath(dirname(dirname(pathof(@__MODULE__))), xs...))
 
+# patch to permutedims
+using Base.Cartesian
+using Base: size_to_strides, checkdims_perm
+for (V, PT, BT) in Any[((:N,), BitArray, BitArray), ((:T,:N), Array, StridedArray)]
+    @eval @generated function Base.permutedims!(P::$PT{$(V...)}, B::$BT{$(V...)}, perm) where $(V...)
+        quote
+            checkdims_perm(P, B, perm)
+
+            #calculates all the strides
+            native_strides = size_to_strides(1, size(B)...)
+            strides_1 = 0
+            @nexprs $N d->(strides_{d+1} = native_strides[perm[d]])
+
+            #Creates offset, because indexing starts at 1
+            offset = 1 - sum(@ntuple $N d->strides_{d+1})
+
+            sumc = 0
+            ind = 1
+            @nexprs 1 d->(counts_{$N+1} = strides_{$N+1}) # a trick to set counts_($N+1)
+            @nloops($N, i, P,
+                    d->(df_d=i_d*strides_{d+1} ;sumc += df_d), # PRE
+                    d->(sumc -= df_d), # POST
+                    begin # BODY
+                        @inbounds P[ind] = B[sumc+offset]
+                        ind += 1
+                    end)
+
+            return P
+        end
+    end
+end
+
 include("bitvector.jl")
 include("arithematics.jl")
 include("networks.jl")
@@ -26,4 +58,4 @@ function __init__()
     @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
 end
 
-end
+end

src/bounding.jl

@@ -59,7 +59,7 @@ function cached_einsum(code::Int, @nospecialize(xs), size_dict)
 end
 function cached_einsum(code::NestedEinsum, @nospecialize(xs), size_dict)
     caches = [cached_einsum(arg, xs, size_dict) for arg in code.args]
-    y = dynamic_einsum(code.eins, (getfield.(caches, :content)...,); size_info=size_dict)
+    y = code.eins(getfield.(caches, :content)...; size_info=size_dict)
     CacheTree(y, caches)
 end
 
@@ -142,4 +142,4 @@ function read_config!(code::NestedEinsum, mt, out)
         end
     end
     return out
-end
+end

src/graph_polynomials.jl

@@ -99,7 +99,7 @@ function contractf(f, gp::GraphProblem; usecuda=false)
     if usecuda
         xs = CuArray.(xs)
     end
-    dynamic_einsum(gp.code, xs)
+    gp.code(xs...)
 end
 
 ############### Problem specific implementations ################
@@ -117,22 +117,6 @@ function generate_tensors(fx, gp::Independence)
     end)
 end
 
-function collect_ixs(ne::NestedEinsum)
-    d = collect_ixs!(ne, Dict{Int,Vector{OMEinsum.labeltype(ne.eins)}}())
-    return [d[i] for i=1:length(d)]
-end
-
-function collect_ixs!(ne::NestedEinsum, d::Dict)
-    for i=1:length(ne.args)
-        if ne.args[i] isa Integer
-            d[ne.args[i]] = collect(OMEinsum.getixs(ne.eins)[i])
-        else
-            collect_ixs!(ne.args[i], d)
-        end
-    end
-    return d
-end
-
 function misb(::Type{T}, n::Integer=2) where T
     res = zeros(T, fill(2, n)...)
     res[1] = one(T)

src/networks.jl

@@ -100,6 +100,23 @@ function labels(code::EinTypes)
     return res
 end
 
+collect_ixs(ne::EinCode) = [collect(ix) for ix in getixs(ne)]
+function collect_ixs(ne::NestedEinsum)
+    d = collect_ixs!(ne, Dict{Int,Vector{OMEinsum.labeltype(ne.eins)}}())
+    return [d[i] for i=1:length(d)]
+end
+
+function collect_ixs!(ne::NestedEinsum, d::Dict)
+    for i=1:length(ne.args)
+        if ne.args[i] isa Integer
+            d[ne.args[i]] = collect(OMEinsum.getixs(ne.eins)[i])
+        else
+            collect_ixs!(ne.args[i], d)
+        end
+    end
+    return d
+end
+
 """
     optimize_code(code; optmethod=:kahypar, sc_target=17, max_group_size=40, nrepeat=10, imbalances=0.0:0.001:0.8, βs=0.01:0.05:10.0, ntrials=50, niters=1000, sc_weight=2.0, rw_weight=1.0)
 

src/viz.jl

@@ -65,8 +65,8 @@ function vizeinsum(nodes, edges; config=zeros(Int, length(nodes)), unit=1.0, gra
     zoom_into(img, XMIN, XMAX, YMIN, YMAX; graphsize=graphsize, rescale=rescale)
 end
 
-function vizeinsum(::EinCode{ixs, iy}, locs::AbstractVector{<:Pair}; kwargs...) where {ixs, iy}
-	vizeinsum(ixs, iy, Dict(locs); kwargs...)
+function vizeinsum(code::EinCode, locs::AbstractVector{<:Pair}; kwargs...)
+    vizeinsum(getixs(code), getiy(code), Dict(locs); kwargs...)
 end
 function vizeinsum(code::NestedEinsum, locs::AbstractVector{<:Pair}; kwargs...)
 	vizeinsum(OMEinsum.flatten(code), locs; kwargs...)