[Good First Issue]: Support torch_scatter::segment_mean_csr

Context

Example model to reproduce operation:

from torch_scatter import segment_csr
import torch


class Example(torch.nn.Module):
    def forward(self, src, indptr):
        return segment_csr(src, indptr, reduce="mean")

src = torch.randn(10, 6, 64)
indptr = torch.tensor([0, 2, 5, 6])
indptr = indptr.view(1, -1)

m = Example()
t = torch.jit.trace(m, example_inputs=(src, indptr))
print(t.inlined_graph)

graph(%self : __torch__.Example,
      %src : Float(10, 6, 64, strides=[384, 64, 1], requires_grad=0, device=cpu),
      %indptr : Long(1, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %3 : NoneType = prim::Constant()
  %4 : Float(10, 3, 64, strides=[192, 64, 1], requires_grad=0, device=cpu) = torch_scatter::segment_mean_csr(%src, %indptr, %3) # /home/vmaxim/openvino/ov_env/lib/python3.10/site-packages/torch/_ops.py:1123:0
  return (%4)

What needs to be done?

As in the main issue.

Example Pull Requests

No response

Resources

• torch_scatter.segment_csr
• Contribution guide - start here!
• Intel DevHub Discord channel - engage in discussions, ask questions and talk to OpenVINO developers
• How to link your Pull Request to an issue

Contact points

As in the main issue.

Ticket

No response