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.

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

Add a Python alternative to seq2seq.gather_tree #1925

Merged
merged 4 commits into from
Jun 21, 2020
Merged

Add a Python alternative to seq2seq.gather_tree #1925

merged 4 commits into from
Jun 21, 2020

Conversation

guillaumekln
Copy link
Contributor

@guillaumekln guillaumekln commented Jun 12, 2020

This PR adds a pure Python TensorFlow implementation of tfa.seq2seq.gather_tree, which can be enabled with the global flag TF_ADDONS_PY_OP. This is useful when Addons custom ops are not readily available (e.g. TensorFlow Serving).

I tested the performance on a real world application: beam search decoding of a neural machine translation model. I'm using a custom beam search implementation but it is close to the BeamSearchDecoder included in Addons.

I did not find significant performance impact. And I think this makes sense: the function is only used at the very end of the decoding and does not involve very complex ops.

@guillaumekln guillaumekln marked this pull request as draft June 15, 2020 08:38
@guillaumekln
Copy link
Contributor Author

Marking this PR as draft. The fixture from #1929 is needed to run tests for both the custom op and the py op.

@guillaumekln guillaumekln marked this pull request as ready for review June 15, 2020 14:00
@guillaumekln
Copy link
Contributor Author

Tests are now running for this new implementation.

Copy link
Member

@qlzh727 qlzh727 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the change. I didn't verify the detail of the beam search implementation, since the newly added test should ensure that the py implementation will have the same behavior as the fused gpu kernel.

Copy link
Member

@seanpmorgan seanpmorgan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM Thx!

@seanpmorgan seanpmorgan merged commit c946906 into tensorflow:master Jun 21, 2020
@guillaumekln guillaumekln deleted the py-gather-tree branch June 29, 2020 12:19
ashutosh1919 pushed a commit to ashutosh1919/addons that referenced this pull request Jul 12, 2020
* Add a Python alternative to seq2seq.gather_tree
* Enable tests for the Python op
jrruijli pushed a commit to jrruijli/addons that referenced this pull request Dec 23, 2020
* Add a Python alternative to seq2seq.gather_tree
* Enable tests for the Python op
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants