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:Yong Tang <yong.tang.github@outlook.com> * Address review comments Signed-off-by:
Yong Tang <yong.tang.github@outlook.com> * Add HostMemory to InTopK kernel Add HostMemory to InTopK kernel based on the review feedback Signed-off-by:
Yong Tang <yong.tang.github@outlook.com>
Loading
Please sign in to comment