How to do a step by step contraction of a tensor network in one cuquantum.Network context ?

[benjaminlanthier]

If I have built a tensor network using the cuQuantum.Network(expr, *operands) class and use this network as a context, would it be possible for me to go through each contraction steps one by one in this context?

I tried building a function update_expr_and_tensors() that does this contraction explicitly, but it is not efficient since it needs to generate a new context each time. That's why I would like to know if there is a built-in process to do that.

The function is the following :

def update_expr_and_tensors(path: list, tensors_list: list, old_expr: str, new_ids: str):
   """
   Purpose:
       Update the 'expr' and the 'tensors_list' in the tensor network from
       the first path step evaluated with cuquantum.Network().contract_path.

   Inputs:
       * path (list of tuples): The optimized path given by cuquantum.Network.contract_path().
       * tensors_list (list of arrays): List of all the tensors represented
       in the tensor network.
       * old_expr (string): Previous 'expr' string given at the tensor network.
       * new_ids (string): Einsum ids for the new tensor.

   Outputs:
       * new_expr (string): The new 'expr' string after only one contraction
       in the network.
       * new_tensors_list (list of arrays): Updated tensors list after the
       contraction of 2 previous tensors.
       * new_info (Optimizer information): Information on the contraction of
       the tensor network.
   """

   # Get the ids of the tensors that will be contracted
   tensors_contracted = path[0] # First contraction done by the path
   t1id = tensors_contracted[0] # id number of the first tensor
   t2id = tensors_contracted[1] # id number of the second tensor

   # Update the old_expr_list to remove the contracted tensors ids and add the new one at the end
   old_expr_list = old_expr.split(',')
   str_id1_to_remove = old_expr_list[t1id]
   str_id2_to_remove = old_expr_list[t2id]
   old_expr_list.remove(str_id1_to_remove); old_expr_list.remove(str_id2_to_remove)
   old_expr_list.append(new_ids)
   new_expr = ','.join(old_expr_list)

   # Get the tensors that will be contracted and contract them to obtain the new tensor
   ti, tj = tensors_list[t1id], tensors_list[t2id]
   subscripts = ','.join([str_id1_to_remove, str_id2_to_remove]) + '->' + new_ids
   new_tensor = cuquantum.contract(subscripts, *[ti, tj])

   # Remove the contracted tensors from 'tensors_list' and add the new one at the end
   tensors_to_remove = [t1id, t2id]
   new_tensors_list = [tensor for i, tensor in enumerate(tensors_list) if i not in tensors_to_remove]
   new_tensors_list.append(new_tensor)

   # Find the new info from those updated inputs
   with cuquantum.Network(new_expr, *new_tensors_list) as tn:
       _, new_info = tn.contract_path({'samples': 500}) # Why specify {'samples': 500} ?

   return new_expr, new_tensors_list, new_info

[yangcal]

Hello Benjamin,

Currently Network.contract directly dispatches all pairwise contractions to our C API cutensornetContractSlices altogether so there is no public API that allows user to tune/walk through pairwise contractions. That being said, one may be able to twerk around our public APIs a bit to achieve the goal.

Before we can further assist you, I noticed a couple of things from your script:

1. It seems that inside update_expr_and_tensors, you're taking a full contraction path (let's call it full_path) from cutensornet.contract_path. In each call, the first pairwise contraction is peformed using cuquantum.contract, which gives you the new_network and then, instead of inheriting the subsequent contraction path from full_path, you again tried to optimize the subsequent path with Network.contract_path. In other words, you're not reusing the full_path but try to again do path searching for the new_network. Can you clarify if this is by design (only making use of the first contraction and iteratively optimizing the path)? If so, for network of n tensors, you're performing n-1 path optimizations which could be costly and I don't think it will give you much if any speedup compared with sticking to full_path.
2. sample=500 is just a safe parameter that we set for intermediate or large network. In practice if you have a small network, you can decrease this number.
3. When you're performing the pairwise contraction using cuquantum.contract, I would recommend to reuse our library handle for better performance, eg, check the Note session in our documentation. This should help you reduce the overhead.

[leofang]

I would like to supplement to Yang's response above that when we return a path, we also return slicing information. If you just take the returned path as-is and ignore the slices, it might not be more performant than, say, a naive left-to-right contraction path. (It is entirely problem dependent.) Conversely, if you intend to perform slicing, manual pairwise contraction like what your update_expr_and_tensors does must be executed with great care, if not impossible.

[benjaminlanthier]

Thank you @yangcal for your quick response. To answer your question, I do evaluate a new path each time I call this function and I must admit that it can be pretty costly, so I will remove it and keep the full_path for now. Also, thank you for the explanations about the sample=500 and the use of the library handle.

Also, thank you @leofang for mentioning manual pairwise contraction using slicing. I saw in the documentation that we can call slices= in the cuquantum.Network.contract() function, but I have troubles finding the right type of data to feed in this parameter for specific contractions. In the function description, it mentions the fact that this input can take the form of a Python sequence object (for arbitrary slice IDs), but I am not sure of what it means. At first, I thought it meant for the input to have the form of the path variable (a list of tuples). I then tried to only contract the first step of the network, so I called cuquantum.contract(slices=path[0]) but I get the following error : 
cuquantum.cutensornet.cutensornet.cuTensorNetError: CUTENSORNET_STATUS_INVALID_VALUE

Is it that, when calling cuquantum.Network.contract_path(), I should ask for a certain number of slices if None is given by default? If so, how can I do it?

[yangcal]

Hello Benjamin,

Let me try to clarify a bit on Leo's answer above. When you request a path with path, info = contract_path(*), the path contains only the pairwise contraction path while the info object contains not just path, but also whether slicing was required and if so what modes are sliced in order to fit the entire calculation into memory. What Leo is highlighting is that, if slicing was actually needed (determined by our algorithm) and you dismiss this information (only using path in your function), there could be undesired consequences, eg, less performant or even not fitting into memory at all.

Going back the slices arg in contract and taking this as an example, in this snippet we manually enforced slicing mode e to extent 2 and mode h to extent 1. That means we're creating 8/2 (from e) * 2/1 (from h) = 8 total slices with slice IDs (0, 1, 2,..., 7). The slices arg in contract allows one to specify how much of the total slices you want to contract, eg, slices=(0,1,2,3) will contract the first 4 slices. slices=None by default will contract all slices but advanced users can use this feature for customized parallelization or research purpose.

Also, I would encourage you to take a look at the OptimizerInfo object returned by contract_path. OptimizerInfo.intermediate_modes holds the mode labels for all intermediates following the contraction path. This can give you a rough idea on what happened during the contraction execution if that's the aim of your function.