Contraction path optimization with EinExprs

[mofeing]

This is an ongoing work with @JoeyT1994 to add support for contraction path optimization using EinExprs.jl. It is a lightweight package ONLY for contraction path optimization. Inspired by opt_einsum and cotengra Python packages, but with more effort spent on simplicity and performance (algorithms occupy 10s of lines and are easy to comprehend).

• Fix lingering Julia 1.10 test errors on Ubuntu and macOS (package extension is loading properly but there are some bugs in the code).
• Fix Windows tests by making KaHyPar.jl an optional dependency of EinExprs.jl, since it is not available for Windows so the extension fails to load since it is looking for a library that doesn't exist.
• Add PackageExtensionCompat.jl to make sure the extension will work for Julia versions below 1.9.

EinExprs is not yet in the General registry, so the @bsc-quantic registry has to be added. We are working towards moving it to the General registry for the next release.
EinExprs has been moved to the General registry.

This is a WIP, not yet ready to be merged.

@JoeyT1994 I have a couple of questions:

• EinExprs represents indices using Symbols. Since an Index has no name, it is ok if I use the id?
• How can I obtain the open indices in a ITensorNetwork?
• What is the expected return type of contraction_sequence?

[JoeyT1994]

Thanks for opening this!

• Yes, the id is a unique identifier for each index so this is a good way to keep track of them.

• The function externalinds(tn) returns a vector of all of the open indices in an ITensorNetwork. Although definitely let us know if you see any cases where it doesn't return what you expect.

• contraction_sequence(tn) returns a nested list. In the case of a vector [A,B,C] of ITensors, contraction_sequence([A,B,C]) might return [3, [1,2]] telling you that you should contract tensors 1 and 2 (A and B in the list) first and then contract the result onto the third tensor in the list (C). In the case of tn, an ITensorNetwork, the return is similar but the entries of the nested list correspond to the vertices v of the ITensorNetwork. These are wrapped as Key(v) to avoid ambiguity when you have vertex names which are lists themselves.

[JoeyT1994]

As an example. The following

g = named_grid((2,2))
s = siteinds("S=1/2", g)
ψ = randomITensorNetwork(s; link_space=2)
contraction_sequence(ψ)

returns

2-element Vector{Vector{Key{Tuple{Int64, Int64}}}}:
 [Key((1, 1)), Key((2, 1))]
 [Key((1, 2)), Key((2, 2))]

telling you to first contract the tensors on sites (1,1) and (2,1) together and the tensors on sites (1,2) and (2,2) together, and then contract the result of those two contractions together.

[mtfishman]

Thanks @mofeing, seems like this will be good to have. Currently we use OMEinsumContractionOrders.jl as a backend, how does that compare to EinExprs.jl?

Adding on to Joey's comments:

1. I think this should be designed as a package extension so it is in optional dependency.
2. ITensor indices are also identified by their tags and prime levels along with their ID. One strategy to handle that could be to make a dictionary mapping each ITensor Index in the network to a symbol that can be used by EinExprs.jl. The hash of an ITensor Index is based on the id, tags, and prime level.

[mofeing]

Thanks @mofeing, seems like this will be good to have. Currently we use OMEinsumContractionOrders.jl as a backend, how does that compare to EinExprs.jl?

EinExprs is a package focused only on contraction path optimization (it cannot perform the einsum operation, that is delegated to the users). EinExprs is both a framework for research on contraction path optimization and easy usage by applications.

In comparison with OMEinsumContractionOrders, it's is far more comprehensible (optimizers are written in 10s lines of code), easier to hack and experiment, and performant. The Exhaustive/Optimal optimizer is almost on par with TensorOperations and in some benchmarks, the Greedy optimizer is x150 faster than the one in OMEinsumContractionOrders. In the future, we would like to add parallelization support for large-scale search.

There is also some support for slicing and visualization so users can see where the cost of their tensor networks is going to be spent before the execution.

Adding on to Joey's comments:

1. I think this should be designed as a package extension so it is in optional dependency.

Okay, great. But I see that you use a kind of "dynamic" check for loading packages using !isdefined(@__MODULE__, :OMEinsumContractionOrders). Do you prefer a package extension or this other way?

2. ITensor indices are also identified by their tags and prime levels along with their ID. One strategy to handle that could be to make a dictionary mapping each ITensor Index in the network to a symbol that can be used by EinExprs.jl. The hash of an ITensor Index is based on the id, tags, and prime level.

mmm so 2 indices with the same id may not be connected?

[mtfishman]

The only reason we don't use package extensions right now for OMEinsumContractionOrders is because we added that dependency before that feature was available in Julia and haven't updated it yet. Ideally for that and this PR we would organize any conditionally loaded code as package extensions and use PackageExtensionCompat.jl for backwards compatibility with older Julia versions.

Yeah, ITensor indices can have the same ID but if they have different prime levels or tags they won't compare equal. Distinguishing indices by either IDs, prime levels, or tags can be convenient in different situations (for example we have a convention that operators distinguish bra and ket spaces using prime levels, which is much more convenient than distinguishing them by IDs which are randomly generated).

[mofeing]

ah ok, then should I use hash(index)?

[mtfishman]

hash alone wouldn't be guaranteed to not give collisions (i.e. return the same hash value even if the indices are different), but using a dictionary would help handle hash collisions.

[mofeing]

mmm then the objectid of the index?

[mtfishman]

Maybe, though it seems like there can still be hash collisions: https://docs.julialang.org/en/v1/base/base/#Base.objectid.

[mtfishman]

I'm picturing something like this:

julia> using ITensors

julia> i, j, k = Index.((2, 2, 2))
((dim=2|id=45), (dim=2|id=388), (dim=2|id=489))

julia> symbols = Dict{Index,Symbol}()
Dict{Index, Symbol}()

julia> using Random

julia> get!(symbols, i, Symbol(randstring()))
:maCJxbai

julia> get!(symbols, j, Symbol(randstring()))
:bfRnBZjQ

julia> get!(symbols, k, Symbol(randstring()))
:wiM0i2Oe

julia> get!(symbols, i, Symbol(randstring()))
:maCJxbai

julia> symbols
Dict{Index, Symbol} with 3 entries:
  (dim=2|id=489) => :wiM0i2Oe
  (dim=2|id=388) => :bfRnBZjQ
  (dim=2|id=45)  => :maCJxbai

Another option is to make a unique vector of all ITensor indices (say with unique) and then use symbols that correspond to each index deterministically based on the position in the vector, i.e.:

julia> inds = unique([i, j, k, k])
3-element Vector{Index{Int64}}:
 (dim=2|id=45)
 (dim=2|id=388)
 (dim=2|id=489)

julia> symbols = Symbol.(eachindex(inds))
3-element Vector{Symbol}:
 Symbol("1")
 Symbol("2")
 Symbol("3")

julia> inds[parse(Int, String(symbols[2]))]
(dim=2|id=388)

[mofeing]

okay, will try to refactor based on your examples. what should I use to get the full list of indices of the TN?

[JoeyT1994]

collect(vertex_data(siteinds(ψ))) will give you the open indices.
collect(edge_data(linkinds(ψ))) will give you the internal indices.
So the concatenation of those two lists should give you them all. @mtfishman maybe there is a function to do this explicitly but I can't see it.

[mtfishman]

@JoeyT1994 does that make a flat list of indices (Vector{<:Index}), or a nested list Vector{<:Vector{<:Index}}? I would have guessed that latter, you probably also need to flatten it with something like collect(Iterators.flatten(v)) or reduce(vcat, v).

I think that is probably the most succinct way for now.

[mtfishman]

Actually I see internalinds here: https://github.com/mtfishman/ITensorNetworks.jl/blob/v0.3.10/src/abstractitensornetwork.jl#L261-L267. Probably the way those functions are implemented should give you an idea for our current best practice for getting lists of indices like that. I think vertex_data and edge_data should be considered more internal functions that most users shouldn't be using (i.e. depend too much on the internal details of the ITensorNetwork data structure).

[mtfishman]

Though the definition of externalinds doesn't look correct to me, I think it should be looping over the vertices instead of the edges.

[JoeyT1994]

@mtfishman #123 should hopefully fix that and export those functions?

[mofeing]

I parametrized the label type in EinExprs v0.6 so now you can use your own Index type. I've updated the code. Can you give it a check @JoeyT1994?

[JoeyT1994]

Thanks @mofeing ! I think we could follow the pattern we have for OMEinsumContractionOrders here and put this code in a file src/requires/EinExprs.jl

Then you can write a function in src/contraction_sequences which checks the EinExprs.jl module exist and then passes to this function ?

That way we don't need EinExprs.jl in the dependencies at all and can just add it to src/test/Project.toml as a required dependency for testing.

[JoeyT1994]

Also could you then add to test/test_contraction_sequence.jl sequences calculated using this new backend?

[mofeing]

Thanks @mofeing ! I think we could follow the pattern we have for OMEinsumContractionOrders here and put this code in a file src/requires/EinExprs.jl

Then you can write a function in src/contraction_sequences which checks the EinExprs.jl module exist and then passes to this function ?

That way we don't need EinExprs.jl in the dependencies at all and can just add it to src/test/Project.toml as a required dependency for testing.

From #120 (comment), I understood you wanted it as a package extension. Currently it's set as a weakdep.

[JoeyT1994]

Ah yes good point. @mtfishman what do you think? Shall we have this as a package extension as it is currently? And then in a future PR also switch OMEEinsum to a package extension like @mofeing has done.

[mtfishman]

The error that CI is throwing now is:

ERROR: LoadError: UndefVarError: `@Algorithm_str` not defined
Stacktrace:
 [1] top-level scope
   @ :0
 [2] include
   @ Base ./Base.jl:495 [inlined]
 [3] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::String)
   @ Base ./loading.jl:2216
 [4] top-level scope
   @ stdin:3
in expression starting at /home/runner/work/ITensorNetworks.jl/ITensorNetworks.jl/ext/ITensorNetworksEinExprsExt/ITensorNetworksEinExprsExt.jl:36
in expression starting at /home/runner/work/ITensorNetworks.jl/ITensorNetworks.jl/ext/ITensorNetworksEinExprsExt/ITensorNetworksEinExprsExt.jl:1
in expression starting at stdin:3
#...

since @Algorithm_str isn't being loaded into the namespace, I committed an attempted fix.

[mtfishman]

@mofeing at least now ITensorNetworksEinExprsExt is getting loaded properly on Ubuntu and macOS but there are some bugs introduced by the refactor, could you take a look?

[mofeing]

I moved KaHyPar to a package extension and bumped the patch version.
JuliaRegistries/General#100771

[mtfishman]

@mofeing still more test failures due to bugs caused by the refactor.

[mtfishman]

I moved KaHyPar to a package extension and bumped the patch version. JuliaRegistries/General#100771

Great, thanks. I think that is a good move.

[mofeing]

@mofeing still more test failures due to bugs caused by the refactor.

The error in "Julia 1" should be solved, but I can't figure out why this error is happening in "Julia 1.7".

https://github.com/mtfishman/ITensorNetworks.jl/actions/runs/7887034139/job/21524035311

[mtfishman]

@mofeing still more test failures due to bugs caused by the refactor.

The error in "Julia 1" should be solved, but I can't figure out why this error is happening in "Julia 1.7".

https://github.com/mtfishman/ITensorNetworks.jl/actions/runs/7887034139/job/21524035311

From what I can see you haven't added and set up PackageExtensionCompat.jl in this PR so the package extension won't work on Julia versions below 1.9.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (e1abe27) 72.73% compared to head (3ec86f2) 73.43%.
Report is 3 commits behind head on main.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #120      +/-   ##
==========================================
+ Coverage   72.73%   73.43%   +0.69%     
==========================================
  Files          68       69       +1     
  Lines        3947     4062     +115     
==========================================
+ Hits         2871     2983     +112     
- Misses       1076     1079       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mofeing]

From what I can see you haven't added and set up PackageExtensionCompat.jl in this PR so the package extension won't work on Julia versions below 1.9.

Ahh, I misunderstood you. I thought you refer to EinExprs, not to this package extension.

Now should be fixed.

[mtfishman]

From what I can see you haven't added and set up PackageExtensionCompat.jl in this PR so the package extension won't work on Julia versions below 1.9.

Ahh, I misunderstood you. I thought you refer to EinExprs, not to this package extension.

Right, both ITensorNetworks.jl and EinExprs.jl need it in order for the extensions to work on older Julia versions.

[mtfishman]

The CI failures on Julia 1.7 are the same ones on main and unrelated to this PR, we'll have to investigate those in a separate PR. Looks like this PR is ready to go, thanks for sticking with all of the comments @mofeing. This will be nice functionality to have, looking forward to testing it out.

[JoeyT1994]

Thanks @mtfishman for the reviewing and @mofeing for the PR. Lets set up a call to discuss running some timing tests to compare einexprs.jl to the other sequence backends