[MXNET-17951] Argsort for fp16

[andrei5055]

Make TopK functions working with fp16 arguments

1. TopK functions (sort, argsort) can work with the fp16 tensors.
2. On GPU the switch from cub::DeviceRadixSort to cub::DeviceSegmentedRadixSort reduces tree mxnet::op::SortByKey calls to one such call in TopKSort function.

Checklist

Essentials

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Code is well-documented:
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Testing TopK functions wit the fp16 tensors ids done
  a) in /tests/python/unittest/test_operator.py:

def test_order():
. . .
    for dtype in [np.float16, np.float32, np.float64]:

b) in /tests/python/unittest/test_ndarray.py

def test_order():
. . .
    tested_types = [np.int32, np.int64, np.float16, np.float32, np.float64]
    for dtype in tested_types:

[mxnet-bot]

Hey @andrei5055 , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [windows-cpu, windows-gpu, edge, unix-gpu, centos-cpu, miscellaneous, unix-cpu, clang, centos-gpu, website, sanity]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[szha]

lint needs fixing. you can test it locally with ci/build.py -R --docker-registry mxnetci --platform ubuntu_cpu --docker-build-retries 3 --shm-size 500m /work/runtime_functions.sh sanity_check (docker and docker-compose required)

[andrei5055]

@mxnet-bot run ci [unix-cpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-cpu]

[andrei5055]

@mxnet-bot run ci [unix-cpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-cpu]

[andrei5055]

@mxnet-bot run ci [unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu]

[andrei5055]

@mxnet-bot run ci [unix-cpu, miscellaneous]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-cpu, miscellaneous]

[andrei5055]

@mxnet-bot run ci [unix-cpu]

[andrei5055]

@mxnet-bot run ci [unix-cpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-cpu]

[szha]

@ptrendx would you take a look?