在Tensorflow中有没有一种与梯度带兼容的求解稀疏线性方程组的有效方法？

Question

我需要以A（y）x=b的形式求解Tensorflow中的方程，其中A是一个大的稀疏带矩阵，也是其他张量（比如y）的函数。当然，解x也是张量y的函数。解完x后，我想取x相对于y的梯度

我考虑了两种选择： 1.使用稀疏外部库有效地反转a，例如scipy.sparse。为此，我需要将张量转换为numpy数组，然后再转换回张量。这种方法的问题是，我不能将渐变磁带与外部库（如scipy.sparse）一起使用。 2.使用Tensorflow的矩阵求逆，它与梯度带一起工作。这对于大型矩阵来说非常缓慢，因为它没有利用张量的稀疏性。我无法在Tensorflow中找到稀疏反转实现

我需要的一个小的简化示例：

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)

当然，A对y的依赖性要比这个例子复杂得多。 在Tensorflow中有什么方法可以做到这一点，或者我必须将自己限制在tf.linalg.inv之内吗