[torch][segment_reduce] Support for multi dimension (cpu only) (#59951)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59951
Add support for multi-d input for cpu forward/backward implementation.
Next step: Adding cuda support for multi-d input.
Test Plan: Added unit tests.
Reviewed By: ngimel
Differential Revision: D29105457
fbshipit-source-id: a389ba4cc10f02434a336b8e7d36259f32552e11
