Need GPU implementation of sparse_segment_* .

[2sin18]

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 sparsesegment* 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]

@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]

@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.