NVIDIA/tensorflow

Need GPU implementation of sparse_segment_* .

2sin18 opened this issue · 2 comments

2sin18 commented

System information

  • TensorFlow version (you are using): r1.15.5+nv22.02
  • Are you willing to contribute it (Yes/No): No

Describe the feature and the current behavior/state.

sparse_segment_sum and other sparse_segment_* ops are important for deep learning models using embedding_lookup. In latest Tensorflow, These ops already has GPU kernels while NVTF still only has CPU kernels

Will this change the current api? How?

No

Who will benefit with this feature?

deep learning models using embedding_lookup.

Any Other info.

No

DEKHTIARJonathan commented

@2sin18 NVTF is TF 1.15, upstream master branch that you point is TF 2.9. If you need these kernels NVIDIA provide TF 2.X containers.

2sin18 commented

@2sin18 NVTF is TF 1.15, upstream master branch that you point is TF 2.9. If you need these kernels NVIDIA provide TF 2.X containers.

Thanks for your quick response. Another 3 questions:

  1. Would NVIDIA open source NVIDIA-TF 2.X too in near future?
  2. My pipelines and libraries still heavily rely on TF 1.x API. What's the future plan of NV-TF 1.x ?
  3. Would NV-TF 1.x consider to accept PR that introduces new op kernels of existed ops? Such optimizations seem do not break compatibility.