Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Frontend][Tensorflow] Sparse dense matmul adjoint option added #7267

Merged
merged 4 commits into from Jan 28, 2021
Merged

[Frontend][Tensorflow] Sparse dense matmul adjoint option added #7267

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2c10b09
[Frontend][Tensorflow] Sparse dense matmul adjoint option added
Jan 13, 2021
8c3c5dc
[1] Review comments handled
Jan 22, 2021
47d211e
[2] Review comments handled
Jan 24, 2021
6d26ce5
[3] Review comments handled
Jan 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
20 changes: 4 additions & 16 deletions python/tvm/relay/frontend/tensorflow.py
View file
Open in desktop
Expand Up @@ -942,7 +942,10 @@ def _impl(inputs, attr, params, mod):
)

if sparse_lhs:
data = _op.transpose(data)
if attr.get("adjoint_a"):
ANSHUMAN87 marked this conversation as resolved.
Show resolved Hide resolved
weight_sp = csr_matrix(weight_sp.transpose())
if not attr.get("adjoint_b"):
data = _op.transpose(data)
else:
weight_sp = csr_matrix(weight_sp.transpose())
ANSHUMAN87 marked this conversation as resolved.
Show resolved Hide resolved

Expand All @@ -955,21 +958,6 @@ def _impl(inputs, attr, params, mod):
if not sparse_lhs:
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
Expand Down
12 changes: 7 additions & 5 deletions tests/python/frontend/tensorflow/test_forward.py
View file
Open in desktop
Expand Up @@ -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(
Expand Down