Derivatives and environments

[amilsted]

If we define a tensor network for e.g. a scalar value (no dangling edges) and we want to compute the derivative of that number with respect to a tensor T in the network, one could use autodiff (depending on the backend). This is also known as computing the "environment" of T.

However, the first derivative is also given by the contraction of the same network with the T tensor removed. For doing this, it would be nice to have a remove_node() method that deletes a node from the network. Any dangling edges attached to that node are also removed. Any connected edges become dangling edges (modified in place, so that the edge objects persist).

This would be useful for a number of algorithms: One can define the network for, say, the energy of a quantum state once, then compute all required environments/derivatives using remove_node(). These are then used to minimize the energy.

Going further, if multiple environments of the same network are desired, it is possible to take the optimal contraction order for one environment and derive optimal contraction orders for all the others: https://arxiv.org/abs/1310.8023
Intermediate results can also often be reused for multiple environments.

One might imagine having a method environments([n1, n2, n3]) (where n1,n2,n3 are nodes) that efficiently computes environments of multiple tensors this way.

Thoughts?

[amilsted]

By the way, I suspect (although I have not checked) that using the proposed environments() has the same asymptotic computational time scaling as using tf.gradients(). Hence it might only be advantageous for backends without autodiff functionality. On the other hand, environments() might be easier to fit into any upcoming automatic multi-device support?

[mganahl]

I think remove_node() is a great idea. It would reduce the length of MERA code (and also others) a lot. It would be good to also have the automated optimal contractor then. @Thenerdstation what is the status on that? I saw the stochastic contractor. Is someone working on a deterministic optimal cotractors?

[chaserileyroberts]

Adam is working on it, but he is a 20%er. I've asked him to add a branch of his work so far, so hopefully that'll be added soon.

[chaserileyroberts]

Remove node should be easy enough to add. Though I think we'll use disconnect to break the edges rather than modify the edges in place. This is consistent with the rest of the code base.

[amilsted]

@Thenerdstation The only trouble I see with replacing the connected edges is that it makes it tricky to keep track of them across the remove_node(). In particular, I would want to know which of the new dangling edges was connected to which axis of the removed node. Ideally, I would not have to remember the axis ordering of the removed node for this to work - the axis names would be enough.

[chaserileyroberts]

Could you give a tiny example of what your ideal workflow would look like? I still don't quite see how modifying the edges in place is more beneficial than just replacing them.

[amilsted]

Something like:

net = TensorNetwork()
n1 = net.add_node(t1, axis_names=['a', 'b', 'c'])
n2 = net.add_node(t2, ...)
...
# (connect edges so that there are no danglings)
...
output_edges = [n1['a'], n1['b'], n1['c']]
net.remove_node(n1)

net.contract_all_naively()  # or whatever :)

env = net.get_final_node()
env.reorder_edges(output_edges)  # want my output_edges to still be part of the network

# I might also want to contract the environment with another node.
# The following replaces the removed node to reproduce the result
# of contracting the original network.
n1 = net.add_node(t1, axis_names=['a', 'b', 'c'])
net.connect(n1['a'], output_edge[0])
net.connect(n1['b'], output_edge[1])
net.connect(n1['c'], output_edge[2])
net.contract_between(n1, env)

If the edges are replaced, one could alternatively have something like
output_edges = net.remove_node(n1, output_edge_axes=['a', 'b', 'c']), I suppose.

[chaserileyroberts]

Oh I see what you mean now. You will need the broken edges for reshaping and to possibly connect it to a different node.

Let me think about this some more. I wanna keep the API as clean and intuitive as possible.

[chaserileyroberts]

Alright this is my compromise. remove_node() will return a Dict[Union[int, Text], Edge] which maps the index/axis name to the newly broken edge. That way, we can use disconnect and keep our current edge pattern consistent and you can continue to use the edge mapping that node had before.

[chaserileyroberts]

@amilsted Please comment on the PR if you have any concerns about the design choice.

[chaserileyroberts]

@mganahl You too.