Reusing contraction paths

[stavros11]

Since many contractors are based on opt_einsum it might be useful to add support for its path reuse functionality. This would save time when contracting the same network multiple times (for example during optimization).

[chaserileyroberts]

Good idea! This is also similar to https://github.com/google/TensorNetwork/issues/73, which we wanted to build for path reuse and parallel/distributed contraction.

[MichaelMarien]

How'd you image the API? Explicit like,

resulting_net, optimal_path= path_algorithm(net) ... path_algorithm(another_net, optimal_path)?

That would be straightforward to implement imo, if @stavros11 has no time, I could take a look for sure.

[stavros11]

This API would work at least for the optimization example. This would look like:

tensors = ... # initialize tensors
net = create_network(tensors)
optimal_path = None
for _ in range(num_epochs):
  result, optimal_path = path_algorithm(net, optimal_path)
  update(tensors) # update tensors using result, gradients, etc.
  net = create_network(tensors)

So path_algorithm should calculate the optimal path only when it is None, otherwise it should use the given path. I don't have problem with time but feel free to take a look if you like, if other people also agree with this API.

Another concern I have is that the optimal path returned by opt_einsum assumes an ordering of nodes in the network, while TensorNetwork stores them in a set which is unordered. If the TensorNetwork ordering changes every time you create a new network (I am not sure if it does), then this strategy won't work.

[chaserileyroberts]

Let me think about the API a little bit longer.

Originally we planned to have ContractionTree proto and ContractionNetwork objects to do something similar. So that would look like this:

# ContractionNetwork is exactly the same as a `TensorNetwork`, except without
# concrete tensors, only `ShapeTensor`s.
net = build_network(...)
shell_network = tensornetwork.ContractionNetwork(net)
# Since ContractionNetwork has the same API as TensorNetwork,
# this should just work.
shell_network = path_algorithm(shell_network)
# The ContractionNetwork will keep track of the ContractionTree,
# which will be implemented as a proto, and thus allows easy
# saving/loading from disk 
save_tree(shell_network.tree, "/path/to/file.pb")
# Networks that are created in the same way as the previous
# network can use the same tree for contraction.
saved_contraction_tree = load("/path/to/file.pb")
new_net = build_network(...)
# And we can contract the new network with the old ContractionTree definition.
new_net = tensornetwork.contract_from_tree_def(new_net, saved_contraction_tree)

I'm not sure this is the BEST way of doing it, but it seems pretty nice. (of course the function names are likely to change).

This also allows for easy distribution of computation, since each branch of the contraction tree can run on independent workers, minimizing communication until absolutely necessary. And that API would just look like

net = build_network(...)
shell_network = tensornetwork.ContractionNetwork(net)
shell_network = path_algorithm(shell_network)
workers = ["list", "of", "remote", "workers", ...]
net = tensornetwork.async_contraction(net, shell_network.tree, workers).

[chaserileyroberts]

Another concern I have is that the optimal path returned by opt_einsum assumes an ordering of nodes in the network, while TensorNetwork stores them in a set which is unordered. If the TensorNetwork ordering changes every time you create a new network (I am not sure if it does), then this strategy won't work.

Nodes have signatures based on the order they are created in the network. As long as you use the same code to contract the network, it should be fine.

[MichaelMarien]

Let me think about the API a little bit longer.

@Thenerdstation ,sure just give a bump here when you settle on the API, I can take it up.

[amilsted]

Just pointing out here that one can also use graph-building or JIT to make path finding a one-time cost.

That said, I think it would be great to have ContractionTree objects, which could also form the basis for automatic parallelization of contractions, efficient computation of multiple derivatives (environments) of tensor networks, and would also be handy for querying the contraction produced by a Contractor (say, to find out the order used, or the cost of each step).