[PT FE] Added support for torch_scatter::segment_mean_csr

[Nishant-ZFYII]

Closes #29734

Details

• Implemented PyTorch Frontend translator for torch_scatter::segment_mean_csr using a prefix-sum + gather boundaries approach (no unrolling over segments).

• The segmentation axis is computed dynamically as axis = indptr.dim() - 1 (matches torch_scatter semantics).

• Added handling for batch/broadcast-style indptr (when axis > 0) via bidirectional broadcast to [src_batch_dims..., indptr_last_dim].

• Uses i64 indices for indptr/shape operations to avoid overflow and to match Gather-supported index types.

• Handles the optional out argument via context.mutate_input(out_idx, result) (if provided).

• Added a regression test for TorchScript that covers:

  • 2D indptr with shape (1, n) (segments along dim 1)
  • 1D indptr (segments along dim 0)

Limitations / Notes

• Requires static ranks for src and indptr (validated).
• Empty segments (indptr[i+1] == indptr[i]) are handled by clamping length to 1 to avoid division by zero (result becomes 0 for empty segments).
• Test is TorchScript-only and is skipped for torch.export / torch.compile backends and when torch_scatter is not installed.

How to test

• PyTorch layer test:

  • python -m pytest tests/layer_tests/pytorch_tests/test_segment_mean_csr.py -k segment_mean_csr -v

Tickets

• [Good First Issue]: Support torch_scatter::segment_mean_csr #29734
• Parent tracking: [Good First Issue]: Operation list to be supported #28584

[Nishant-ZFYII]

@mvafin Can you please review .

Thanks. !

[Nishant-ZFYII]

@mlukasze @p-wysocki

Good Day,

Kindly requesting you to review the PR. Please let me know if further changes are required.

Thanks and Regards.