Eigensolving TN operators

[chaserileyroberts]

We should build a method to allow users to find the top eigenvector of a linear operator defined by a tensor network.

My API proposal is as such

tensor = tn.eigensolution(in_edges=list_of_edges, out_edges=list_of_edges)

Where tensor.shape == [edge.dimension for edge in in_edges] and tensor is the actual tensor and NOT a Node.

You can easily find the necessary nodes via tn.reachable, and to prevent smashing the given tensornetwork, we can safely copy all of the nodes via tn.replicate_nodes.

[chaserileyroberts]

With the new lanczos algorithm written up, this is actually a very viable function now.

[mganahl]

I agree we should make sparse solvers directly accessible for users, not only via the backend object. I see two potential problems with the API above.

First, linear operators come with different properties (hermitian, symmetric, positive definite, ....), and the choice of which algorithm to use has to take that into account. Instead of having a single function eigensolution, I'd rather have different functions like eigsh_lanczos, eigs_arnoldi, eigsh_davidson and so on. The alternative to that would be to pass flags to eigensolution, which typically is a major source of obscure bugs, so I'd like to avoid it as much as possible.

The second point is about passing edges. The problem is that the network behind the passed edges can be arbitrarily complicated. The network contraction needs to be called many times during a single call to eigensolution, so to get good performance we need a good contraction path, e.g. from an automatic contraction path finder (they are very costly though). In a typical setting eigensolution is called many times, e.g. in a for or while loop. To maintain performance, we have to reuse the previously computed contraction paths, and I'm not sure if there is a clean and easy way to keep track of all these paths for an arbitrary number of networks.

I would instead keep the API identical to the one used in the backend, i.e.

tn.eigs_arnoldi(matvec, args_to_matvec, ...)

so that a user can do

#matvec maps an input `vector` to an output vector of same dimension.
def matvec(t1, t2, t3,..., tn, vector):
  """construct the network here 
  """
  return result

eigvals, eigvecs = tn.eigs_arnoldi(matvec,[t1, t2, t3, ..., tn],...)

Contraction path finders can be directly involved by the user by using e.g. ncon

#construct the network
t1[0] ^ t4[2]
t3[4] ^ t2[2]
vector[1] ^ t15[0]
...

contraction_path = find_best_path([t1, t2, ..., tn, vector]) #or use edges instead 
def matvec(t1, t2, t3,..., tn, vector):
  return ncon([t1,t2,...,tn, vector],contraction_path)
eigvals, eigvecs = tn.eigs_arnoldi(matvec,[t1, t2, t3, ..., tn],...)

Note that matvec is not a closure here, so this could potentially lead to issues with jax (other backends are fine). I did some preliminary tests though and it seems that this kind of thing is fine with jax (we need to investigate this though in more detail).

[chaserileyroberts]

First, linear operators come with different properties (hermitian, symmetric, positive definite, ....

That's a really good point. Yeah, lets change the api the be eigensolve_hermitian, eigensolve_symmetric, etc.

The second point is about passing edges. The problem is that the network behind the passed edges can be arbitrarily complicated....

The problem I see with having a user provide their own matvec function is that graphs are mutable, so the second time the matvec function is called, there would be lots of errors. By having the user provide the edges instead, we can copy the graph via tn.reachable(edges) and tn.copy. Then we can safely build our own matvec function that preservers the original network.

So how about this instead?

tn.eigensolve_hermitian(input_edges, output_edges, nodes, contraction_algorithm, *args).

Here, the input_edges and output_edges remain the same, but we add two additional parameters nodes, and contraction_algorithm. If a user wants to hand define a contraction path, then they would pass a contraction_algorithm that's defined like:

def my_path(nodes):
  # The "input vector" will always be appended to the start of the list.
  tmp = nodes[0] @ nodes[1]
  return tmp @ nodes[2]

tn.eigensolution_hermitian(inputs, outputs, my_nodes, my_path)

Then we can take their function and wrapped it safely into our own matvec function. Additionally, a user could freely pass any of the tn.contractors. If no method is passed, we'll default to using auto.

[chaserileyroberts]

nodes would also be an optional argument. If nodes is not passed, we'll default to whatever list(tn.reachable(input_edges + output_edges)) returns

[gevenbly]

Hey guys, I am not so familiar with the details and workings of the tensornetwork library yet, but I thought I could offer some general suggestions.

In practice, one often needs to solve eigenvectors of a linear operator defined by a weighted sum of several different networks (or maybe things even more complicated), so I think that it is essential to allow the user to define their own matvec function.

As for the API, I would suggest sticking with (or adopting something close to) the existing one for sparse eigensolvers: scipy.sparse.linalg.eigsh and scipy.sparse.linalg.eigs, which can be made compatible with an arbitrary matvec function by using scipy.sparse.linalg.LinearOperator or similar. The eigsh and eigs already employ the optimal sparse eigensolver (given the symmetry of the linear operator) and accept a host of optional arguments (k, M, sigma, which, v0, ncv, maxiter, tol, return_eigenvectors) which I feel should be supported (in my own work, at least, I routinely require the use of all of these).

[chaserileyroberts]

Hm, in that case it might be nice to actually just have these in JAX proper. Arbitrary matvec functions seems somewhat out of scope for this library, where as putting eigsh into JAX would be helpful for many more users.

[mganahl]

It would indeed be, as would eigs.

I actually think that arbitrary matvec has to absolutely be within the scope of this library. Sparse solvers are needed in so many tensor network algorithms that not having them would be an absolute no-go (I can say with confidence that I would just not use the library in that case). As @gevenbly has pointed out, not all matvecs can be easily written as a simple tensor network, and we have to support such cases as well. I think the best way of providing this is via user defined matvecs as in the above-illustrated way.

One of the main points of tensornetwork is to provide functionality to write backend independent tensor network algorithms, so we have to provide eigs and eigsh not only in the backend but also via the front end. If we don't provide the matvec fuctionality for eigs and eigsh to the user via the front end, I'm concerned that many will just opt to not use it.

I think to maintain backend independence it makes sense to have eigsh take Node objects and return Node objects.

The good news is that we essentially have almost all the pieces. For user-defined contraction paths this would work fine:

def matvec([node1, node2, ...], vector_node,...):
  #user uses whatever API she prefers:
  #tensornetwork API
  node1[0] ^ node2[5]
  ...
  res = ... #the final node
  #or ncon
  res = tn.ncon([node1, node2, ...], [[...],[...],....])
  return res

eigvals, eignodes = eigsh(matvec, [n1,n2,n3],...) #returns eigenvalues and eigenvectors, wrapped into `Node` objects.

Optimal contraction paths could be implemeted using e.g. ncon, or have matvec take the path as argument. In pytorch and numpy all of this is straightforward and already supported.

JAX is a bit more tricky because we need to jit the code, and passing nodes or other types as arguments to jitted functions is by default not permitted, but we can register user-defined types with JAX, which makes them passable. I did some quick tests and this seems to work with Node objects, so the above API would likely work there as well. The only remaining question mark is tensorflow, as I am not sure if it is possible to pass user-defined arguments to tf.fuction wrapped functions. I'm confident though we can find a solution there as well to get a seamless API for all backends.

I'm open to using an API with a LinearOperator class, or just passing plain functions. One thing we probably need to keep though is passing an args argument containing the input arguments to the matvec operation. JAX strongly recommends jitted functions being closures, otherwise unexpected side effects could appear. Unless I fully understand the implications of passing non-closures I'd rather stick with their recommendation. This essentially would force this part of the API on us.

[mganahl]

In fact, there are even easier workarounds than using the JAX pytree, and those also work with tensorflow

[chaserileyroberts]

Let's just completely ignore tensorflow for this use case, as I'm not convinced scientists will ever use the TF backend, especially since we highly encourage using JAX when you want to use hardware accelerators anyway. TF and PyTorch are mainly supported for ML use cases, which likely would never need eigs.

I think to maintain backend independence it makes sense to have eigsh take Node objects and return Node objects.

I'm not sure that's the best idea. The first problem is what do you do when these nodes are already connected to each other. Raise an error? Assume it's ok? It's far too much of a footgun for users. Like you said, JAX expectes functions to be closures, and the tn.Nodes are highly mutable in our current implementation. Another footgun.

One option to work around this is to instead pass tensors instead of nodes to the matvec function. But if we do that, then when is the user using a tensor network? It feels like its at the wrong level of abstraction.

[mganahl]

Let's just completely ignore tensorflow for this use case, as I'm not convinced scientists will ever use the TF backend, especially since we highly encourage using JAX when you want to use hardware accelerators anyway. TF and PyTorch are mainly supported for ML use cases, which likely would never need eigs.

I think to maintain backend independence it makes sense to have eigsh take Node objects and return Node objects.

I'm not sure that's the best idea. The first problem is what do you do when these nodes are already connected to each other. Raise an error? Assume it's ok? It's far too much of a footgun for users. Like you said, JAX expectes functions to be closures, and the tn.Nodes are highly mutable in our current implementation. Another footgun.

This should clearly raise an error, since all previously defined connections will be erased during the first call to matvec.

If we require input nodes to be unconnected I don't see a problem with the closures. In that case, Nodes can essentially be regarded as equivalent to jax-arrays. Unless the user has decided to do some really shady things with references to Node.tensor, I don't see how this could break.

One option to work around this is to instead pass tensors instead of nodes to the matvec function.

I think this would be possible too.

But if we do that, then when is the user using a tensor network? It feels like its at the wrong level of abstraction.

He uses it exactly where it has always been used: inside the matvec function. I think we should not regard a tensor network as a single entity that is to be defined at one place in the code at the beginning and maintains its connections for eternity. If this is our design philosophy, that would lead to a very brittle API. In practice, every tensor network calculation is broken down into contractions and decompositions over small subsets of tensors. There may be special cases of tensor networks where a global approach could be feasible and convenient, and they are worth exploring. But for the large majority, it will lead to an overly verbose and complicated code.

The small subset of tensors is the entity we should focus on, and for which we should provide backend-independent code to do contractions and dense and sparse decomposition.

[chaserileyroberts]

Alright, you've convinced me.

Let's do the matvec functions with tensors being passed in.

Let's make the API look like this?

def eigs(matvec, tensors, ... <eigs_args>...):
  ...

And then the matvec functions should follow the API spec of

def matvec(vector, tensors):
  ...

My annoyance with this API is that matvec has to be define with tensors, even if they don't use them.

[mganahl]

Sounds good. The current backends expect matvec(tensors, vector), but I'll fix this to conform to the above.

I agree, having Nodes would be the natural choice. This is a matter of making this work with JAX, I'll have a look at this.

[mganahl]

Is there a reason for not using

def matvec(vector, n1, n2, ,...):
  #do stuff

and call matvec(vector, *args) inside eigsh?
this would prevent uses from needing to do this:

def my_matvec(tensor, args):
   t1, t2, t3 = args
   #do stuff with t1 t2, t3

or even worse:

def my_matvec(tensor, args):
  args[0][0] ^ args[1][2]
  args[1][5] ^ args[2][3]
  ...

It's a minor thing

[gevenbly]

The format you suggest:
def matvec(vector, n1, n2, ,...):
is the common one I have seen before (e.g. also inbuilt with MATLAB eigs and whatnot).

In general the arguments n1,n2,... could be of any type, not just a set of tensors. It is common for matvec functions to need other scalar parameters, or to accept custom objects (e.g. for specifying contraction order)

[chaserileyroberts]

Yes def f(vector, *args) sounds good to me too.