Add gather/scatter support 1D tensor

[chang-l]

This PR is to add gather/scatter support 1D tensor on python level, as WholeGraph should support basic indexing operations for both 1D (array) and 2D (matrix) wholememory tensors. Without this PR, if with 1D wholememory tensor, gather/scatter op does not work, e.g.,

wholegraph/python/pylibwholegraph/pylibwholegraph/torch/tensor.py

Line 89 in 0efba33

assert input_tensor.dim() == 2

To test, run

pytest --cache-clear  --import-mode=append  tests/wholegraph_torch/ops/test_wholegraph_gather_scatter.py -s

Remaining issue:

On my local test with single GPU, the test can pass.
For multiGPU setup, gather op works fine, but 1D scatter seems not working as it would crash at:

wholegraph/python/pylibwholegraph/pylibwholegraph/tests/wholegraph_torch/ops/test_wholegraph_gather_scatter.py

Line 108 in 2e963b9

assert torch.allclose(local_tensor, local_tensor_ref)

with incorrect scatter outputs: Indices where allclose fails: tensor([0., 0., 0., ..., 0., 0., 0.]) tensor([ 1435., 1439., 1443., ..., 257703., 257707., 257711.])

@linhu-nv Can you please take a look? Does scatter_op suppose to work with 1D wholememory tensor?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[linhu-nv]

Thanks for bring this up. Yes, scatter_op supposes to work with 1D wholememory tensor. I will try to find out why it does't work.