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
 
@@ -96,7 +96,7 @@ function bounding_contract(@nospecialize(code::EinCode), @nospecialize(xsa), yma
     bounding_contract(NestedEinsum((1:length(xsa)), code), xsa, ymask, xsb; size_info=size_info)
 end
 function bounding_contract(code::NestedEinsum, @nospecialize(xsa), ymask, @nospecialize(xsb); size_info=nothing)
-    size_dict = OMEinsum.get_size_dict(getixs(flatten(code)), xsa, size_info)
+    size_dict = OMEinsum.get_size_dict(getixs(flatten(code)), (xsa...,), size_info)
     # compute intermediate tensors
     @debug "caching einsum..."
     c = cached_einsum(code, xsa, size_dict)
@@ -142,4 +142,4 @@ function read_config!(code::NestedEinsum, mt, out)
         end
     end
     return out
-end
+end

src/graph_polynomials.jl

@@ -97,7 +97,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 ################

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(flatten(code), locs; kwargs...)
@@ -76,4 +76,4 @@ function vizeinsum(ixs::Tuple, iy::Tuple, locs::Dict; kwargs...)
 	nodes = [l=>locs[l] for l in legs]
 	edges = [map(i->findfirst(==(i), legs), ix) for ix in ixs]
 	vizeinsum(nodes, edges; config=[l ∈ iy for l in legs], kwargs...)
-end
+end