[v1.x] fixing breaking change introduced in #17123 when batch_axis=0

[access2rohit]

Description

earlier split_and_save on batch_axis=0 would take in CSRNDArray and return CSRNDArray but now its not doing the same and returns NDArray. This causes change in external API behaviour. More details in #19268. Change was introduced in #17123

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

Behaviour After Fix

import scipy.sparse as sps
import mxnet.ndarray.sparse as mxsps
import mxnet as mx
x = mxsps.csr_matrix(sps.coo_matrix(([2.0], ([99], [999]))).tocsr(), ctx=mx.gpu(0))
y = mx.gluon.utils.split_and_load(x, (mx.gpu(0), mx.gpu(0)))
print(x)
print(y)

output:

[20:54:45] ../src/base.cc:84: Upgrade advisory: this mxnet has been built against cuDNN lib version 7501, which is older than the oldest version tested by CI (7600).  Set MXNET_CUDNN_LIB_CHECKING=0 to quiet this warning.

<CSRNDArray 100x1000 @gpu(0)>
[
<CSRNDArray 50x1000 @gpu(0)>,
<CSRNDArray 50x1000 @gpu(0)>]

[mxnet-bot]

Hey @access2rohit , 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: [clang, edge, centos-cpu, centos-gpu, miscellaneous, website, windows-cpu, unix-cpu, windows-gpu, unix-gpu, 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.

[sandeep-krishnamurthy]

@sxjscience

[sxjscience]

I think we may need to add a test case. Also, we should better add a comment in the code base that links to the issue. The root cause is that slice_axis has not supported sparse ndarray.

[access2rohit]

@mxnet-bot run ci [all]

[mxnet-bot]

Jenkins CI successfully triggered : [website, sanity, centos-gpu, centos-cpu, edge, windows-cpu, unix-cpu, clang, windows-gpu, miscellaneous, unix-gpu]

[access2rohit]

@sxjscience can you review again ?

[sxjscience]

LGTM