This is an automated email from the ASF dual-hosted git repository.
zhaowu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new dda8f5d [Frontend][Tensorflow] Sparse dense matmul adjoint option
added (#7267)
dda8f5d is described below
commit dda8f5d944747b9f48b9155e866fd0f746fcd9bb
Author: ANSHUMAN TRIPATHY <[email protected]>
AuthorDate: Thu Jan 28 11:28:13 2021 +0530
[Frontend][Tensorflow] Sparse dense matmul adjoint option added (#7267)
* [Frontend][Tensorflow] Sparse dense matmul adjoint option added
* [1] Review comments handled
* [2] Review comments handled
* [3] Review comments handled
---
python/tvm/relay/frontend/tensorflow.py | 69 ++++++++++++++++--------
tests/python/frontend/tensorflow/test_forward.py | 12 +++--
2 files changed, 53 insertions(+), 28 deletions(-)
diff --git a/python/tvm/relay/frontend/tensorflow.py
b/python/tvm/relay/frontend/tensorflow.py
index 2c7361a..b34e6c7 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -926,13 +926,6 @@ def _sparse_tensor_dense_matmul():
data = inputs[3]
- # By default, in tensorflow the first input ,i.e., data is sparse
- sparse_lhs = True
-
- # If both are true means First input was dense and second was sparse
- if attr.get("adjoint_a") and attr.get("adjoint_b"):
- sparse_lhs = False
-
rows = [x[0] for x in indices_tensor]
cols = [x[1] for x in indices_tensor]
@@ -941,9 +934,53 @@ def _sparse_tensor_dense_matmul():
(values_tensor, (rows, cols)),
shape=tuple(dense_shape_tensor.tolist())
)
- if sparse_lhs:
+ # As per tensorflow implementation, we have 4 possible input
combination
+ # and the first input(A) is always sparse and second input(B) is
always dense.
+ # Case 1: A , B , adjoint_a=False, adjoint_b=False --> A * B
+ # Case 2: A , B , adjoint_a=True, adjoint_b=False --> A.T * B
+ # Case 3: A , B , adjoint_a=False, adjoint_b=True --> A * B.T
+ # Case 4: A , B , adjoint_a=True, adjoint_b=True --> A.T * B.T
+ #
+ # Topi implementation for sparse_dense(matmul) has 2 possible input
+ # combination where first input(A) is always dense
+ # and second input(B) is always sparse.
+ # Case 1: A , B, sparse_lhs = False --> A * B.T
+ # Case 2: A , B, sparse_lhs = True --> B * A.T
+ #
+ # The mapping would be as below:
+ # TF Case 1: A , B , adjoint_a=False, adjoint_b=False
+ # --> In TF: A * B --> In Topi: A * B.T.T
+ # --> sparse_dense(transpose(B), A, sparse_lhs=True)
+ #
+ # TF Case 2: A , B , adjoint_a=True, adjoint_b=False
+ # --> In TF: A.T * B --> In Topi: A.T * B.T.T
+ # --> sparse_dense(transpose(B), transpose(A),
sparse_lhs=True)
+ #
+ # TF Case 3: A , B , adjoint_a=False, adjoint_b=True
+ # --> In TF: A * B.T --> In Topi: A * B
+ # --> sparse_dense(B, A, sparse_lhs=True)
+ #
+ # TF Case 4: A , B , adjoint_a=True, adjoint_b=True
+ # --> In TF: A.T * B.T --> In Topi: (B * A.T).T
+ # --> transpose(sparse_dense(B, transpose(A),
sparse_lhs=False))
+
+ # By default, in tensorflow the first input ,i.e., data is sparse
+ sparse_lhs = True
+
+ # TF Case 1:
+ if not attr.get("adjoint_a") and not attr.get("adjoint_b"):
+ data = _op.transpose(data)
+ # TF Case 2:
+ elif attr.get("adjoint_a") and not attr.get("adjoint_b"):
data = _op.transpose(data)
+ weight_sp = csr_matrix(weight_sp.transpose())
+ # TF Case 3:
+ elif not attr.get("adjoint_a") and attr.get("adjoint_b"):
+ pass
+ # TF Case 4:
+ # attr.get("adjoint_a") and attr.get("adjoint_b"):
else:
+ sparse_lhs = False
weight_sp = csr_matrix(weight_sp.transpose())
weight_data = _expr.const(weight_sp.data, weight_sp.data.dtype)
@@ -953,23 +990,9 @@ def _sparse_tensor_dense_matmul():
ret = _op.nn.sparse_dense(data, [weight_data, weight_indices,
weight_indptrs], sparse_lhs)
if not sparse_lhs:
+ # TF Case 4
ret = _op.transpose(ret)
- # Case 1. If both are true means first input was dense and second was
sparse
- # Case 2. If both are false means first input was sparse and second
was dense
- # TODO(ANSHUMAN87): Support other adjoint option too
- if not (
- (attr.get("adjoint_a") and attr.get("adjoint_b"))
- or ((not attr.get("adjoint_a")) and (not attr.get("adjoint_b")))
- ):
- raise tvm.error.OpAttributeUnImplemented(
- "Only tf.sparse.sparse_dense_matmul() with adjoint_a=True and
adjoint_b=True"
- "or with adjoint_a=False and adjoint_b=False"
- " is supported, but adjoint_a={} and adjoint_b={} was
supplied.".format(
- attr.get("adjoint_a"), attr.get("adjoint_b")
- )
- )
-
return ret
return _impl
diff --git a/tests/python/frontend/tensorflow/test_forward.py
b/tests/python/frontend/tensorflow/test_forward.py
index 3c30b66..34ee0f3 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -1758,19 +1758,21 @@ def test_forward_batch_matmul():
# ----------------------------------
-def _test_sparse_dense_matmul(indices, values, A_shape, B_shape, dtype,
flip=False):
+def _test_sparse_dense_matmul(indices, values, A_inp_shape, B_inp_shape,
dtype, flip=False):
""" One iteration of sparse_dense_matmul """
- # TODO(ANSHUMAN87): Support adjoint options too
- for adjoint_a in [False]:
- for adjoint_b in [False]:
+ for adjoint_a in [False, True]:
+ for adjoint_b in [False, True]:
+ A_shape = A_inp_shape[::-1] if adjoint_a else A_inp_shape
+ B_shape = B_inp_shape[::-1] if adjoint_b else B_inp_shape
+
with tf.Graph().as_default():
A_sp = tf.sparse.SparseTensor(indices=indices, values=values,
dense_shape=A_shape)
B = tf.placeholder(shape=B_shape, dtype=dtype, name="B")
if flip:
result = tf.sparse.sparse_dense_matmul(
- B, A_sp, adjoint_a=adjoint_a, adjoint_b=adjoint_b
+ B, A_sp, adjoint_a=adjoint_b, adjoint_b=adjoint_a
)
else:
result = tf.sparse.sparse_dense_matmul(
