Sparse matrix triangular solve in Tensorflow?

Question

Has there been an implementation for solving Ax = b with a sparse triangular matrix in Tensorflow? (corresponds to the tf.matrix_triangular_solve())

AFAIK, if we have A, for example, as a lower triangular matrix with a sparse matrix representation, we need to convert it to a dense matrix with the tf.sparse_to_dense().

But, if A has a very large dimension, e.g., 16416x16416, and very sparse entries, e.g., 0.018% (about 45216 non-zeros), it would take a large amount of memory.

I think it would be very helpful if we could take advantage of sparse matrix solver, such as for matrix with a banded structure, in Tensorflow.

Sorry, if my question is not relevant. For example, if there have been any solutions for this, I would be grateful for any help.

Thanks.

user3446746 user3446746 · Accepted Answer · 2017-10-24T14:50:21

I had the same problem and I created a custom op for it. As long as you don't want to compute gradients into A and A remains fixed, then this code should help:

import tensorflow as tf
import numpy as np
from scipy.sparse import linalg as sla
import scipy

lu = sla.splu(A)

# Define custom py_func which takes also a grad op as argument:
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)


def sparse_solve(x, lu, dtype=tf.float64, name=None):

    with tf.name_scope(name, 'SparseSolve', [x]) as name:
        solve_x = py_func(lu.solve,
                        [x],
                        [dtype],
                        name=name,
                        grad=_SparseSolveGrad(dtype, lu))
        return solve_x[0]

class _SparseSolveGrad:
    def __init__(self, dtype, lu):
        self.dtype = dtype
        self.lu = lu

    def __call__(self, op, grad):
        x = op.inputs[0]
        y = tf.conj(tf.py_func(self.lu.solve, [tf.conj(grad)], self.dtype))
        return y

The solution is based on code I found at https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342

At least for me, the solution is very fast. Let me know if you find an error (e.g. in the gradient calculations)

Sparse matrix triangular solve in Tensorflow?

2 Answers