I need to solve equations in Tensorflow in the form A(y)x = b, where A is a large sparse band matrix and also a function of some other tensor say y. Naturally, the solution x will be a function of tensor y too. After solving for x, I want to take gradient of x with respect to y.
I considered two options:
1. Use a sparse external library to efficiently invert A, such as scipy.sparse. For this I need to convert the tensors to numpy array and then back to tensors. The problem with this approach is that I cannot use gradient tape with external libraries such as scipy.sparse.
2. Use Tensorflow's matrix inversion that works with gradient tape. This is extremely slow for large matrices, since it does not utilize the sparsity of the tensor. I was unable to find a sparse invert implementation in Tensorflow.
A small simplified example of what I need:
y = tf.constant(3.14)
A = my_sparse_tensor(shape=(1000, 1000)) # Arbitrary function that returns a sparse tensor
b = tf.ones(shape=(1000, 1))
with tf.GradientTape() as g:
g.watch(y)
A = A * y
x = tf.matmul(sparse_invert(A), b)
dx_dy = g.gradient(x, y)
Of course the dependence of A on y is much more complicated than in this example.
Is there any way to do this in Tensorflow, or do I have to restrict myself to tf.linalg.inv ?