code: def test_softmax():
      src = torch.tensor([1.0, 1.0, 1.0, 1.0])
      index = torch.tensor([0, 0, 1, 2])
      ptr = torch.tensor([0, 2, 3, 4])
  
      out = softmax(src, index)
      assert out.tolist() == [0.5, 0.5, 1, 1]
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert softmax(src, ptr=ptr).tolist() == out.tolist()
      else:
          with pytest.raises(
              NotImplementedError, match="requires 'index'"
          ):
              softmax(src, ptr=ptr)
  
      src = src.view(-1, 1)
      out = softmax(src, index)
      assert out.tolist() == [[0.5], [0.5], [1], [1]]
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert softmax(src, None, ptr).tolist() == out.tolist()
  
      jit = torch.jit.script(softmax)
      assert torch.allclose(jit(src, index), out)
  
imports: import pytest
  import torch
  import torch_geometric.typing
  from torch_geometric.profile import benchmark
  from torch_geometric.utils import softmax
label_desc: Use image node_id 1 for calling a global function with example usage: test_softmax() without return types