-
Notifications
You must be signed in to change notification settings - Fork 3
/
tf_sparse_utils.py
73 lines (57 loc) · 2.66 KB
/
tf_sparse_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import tensorflow as tf
import numpy as np
from scipy.sparse import coo_matrix
def to_sparse(a):
# tf.contrib.layers.dense_to_sparse has a memory leak issue, see:
# https://github.com/tensorflow/tensorflow/issues/17590
zero = tf.constant(0, a.dtype)
indices = tf.where(tf.not_equal(a, zero))
values = tf.gather_nd(a, indices)
shape = a.shape
return tf.SparseTensor(indices, values, shape)
def to_dense(a):
# the tf.sparse.to_dense function does not support back-propagation, see:
# https://github.com/tensorflow/tensorflow/issues/6391
# bug is fixed in Tensorflow 1.14: https://github.com/tensorflow/tensorflow/releases/tag/v1.14.0
# for older versions, workaround is performing addition with dense zero tensor
# return tf.sparse.add(a, tf.zeros(a.shape, dtype=a.dtype))
return tf.sparse.to_dense(a)
def csr_to_sparse_tensor(a):
# https://stackoverflow.com/a/43130239
a_coo = a.tocoo()
a_coo.eliminate_zeros()
indices = np.mat([a_coo.row, a_coo.col]).transpose()
tensor = tf.SparseTensor(indices, a_coo.data, a_coo.shape)
return tf.sparse.reorder(tensor)
# return tensor
def sparse_tensor_to_csr(a):
indices = a.indices.numpy()
rows = indices[:, 0]
cols = indices[:, 1]
data = a.values.numpy()
shape = (a.shape[0].value, a.shape[1].value)
a_coo = coo_matrix((data, (rows, cols)), shape=shape)
return a_coo.tocsr()
def sparse_multiply(a, b):
# there are multiple ways to do it in Tensorflow, see:
# https://www.tensorflow.org/api_docs/python/tf/sparse/sparse_dense_matmul
# https://stackoverflow.com/questions/34030140/is-sparse-tensor-multiplication-implemented-in-tensorflow
dense_a = to_dense(a)
dense_b = to_dense(b)
return tf.matmul(dense_a, dense_b, a_is_sparse=True, b_is_sparse=True)
def pad_diagonal(a, padded_length):
# given square matrix "a", pad with 1's on diagonal until matrix is of size "padded_length"
a_length = a.shape[0].value
if a_length > padded_length:
raise RuntimeError(f"padded length {padded_length} is larger than matrix length {a_length}")
if a_length == padded_length:
return a
a_indices = a.indices
new_range = tf.range(a_length, padded_length, dtype=tf.int64)
new_indices = tf.tile(tf.expand_dims(new_range, axis=1), [1, 2])
padded_indices = tf.concat([a_indices, new_indices], axis=0)
a_values = a.values
new_values = tf.ones(padded_length - a_length, dtype=tf.float64)
padded_values = tf.concat([a_values, new_values], axis=0)
padded_shape = [padded_length, padded_length]
return tf.SparseTensor(indices=padded_indices, values=padded_values, dense_shape=padded_shape)