Commit 4d5ff2ee authored by Yong Tang's avatar Yong Tang Committed by drpngx
Browse files

Support placeholder for parameter k in tf.nn.in_top_k (#11197)



* Support placeholder for parameter k in tf.nn.in_top_k

This fix tries to address the issue raised in #9717 where
it was not possible to have tensor for k in nn.in_top_k.

This fix adds the implementation of InTopKV2Op, adds addition test
cases, and following similiar workflow in #10840:
1. Register new kennel InTopKV2Op
2. Hide InTopK and InTopKV2 in python (tensorflow/python/ops/hidden_ops.txt)
3. Add a wrapper in_top_k (in tensorflow/python/ops/nn_ops.py) pointing to gen_nn_ops._in_top_k

Another PR will be created after 3 weeks once this PR is merged:
1. Change the implementation of the wrapper in_top_k (in tensorflow/python/ops/nn_ops.py)
pointing to gen_nn_ops._in_top_kv2

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Address review comments

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>

* Add HostMemory to InTopK kernel

Add HostMemory to InTopK kernel based on the review feedback

Signed-off-by: default avatarYong Tang <yong.tang.github@outlook.com>
parent bad0e44c
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment