Revert "Support multi-dimensional lengths in segment_reduce to support pytorch_scatter.segment_* functionalities (CUDA)"
This reverts commit 40f7ef1f3db9717d8149a0bd1e8b8c80c8600753.
Reverted https://github.com/pytorch/pytorch/pull/77061 on behalf of https://github.com/janeyx99 due to Broke segment_reduce tests on trunk, e.g., https://hud.pytorch.org/pytorch/pytorch/commit/40f7ef1f3db9717d8149a0bd1e8b8c80c8600753