code: def test_softmax_dim():
      index = torch.tensor([0, 0, 0, 0])
      ptr = torch.tensor([0, 4])
  
      src = torch.randn(4)
      assert torch.allclose(
          softmax(src, index, dim=0), src.softmax(dim=0)
      )
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert torch.allclose(
              softmax(src, ptr=ptr, dim=0), src.softmax(dim=0)
          )
  
      src = torch.randn(4, 16)
      assert torch.allclose(
          softmax(src, index, dim=0), src.softmax(dim=0)
      )
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert torch.allclose(
              softmax(src, ptr=ptr, dim=0), src.softmax(dim=0)
          )
  
      src = torch.randn(4, 4)
      assert torch.allclose(
          softmax(src, index, dim=-1), src.softmax(dim=-1)
      )
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert torch.allclose(
              softmax(src, ptr=ptr, dim=-1), src.softmax(-1)
          )
  
      src = torch.randn(4, 4, 16)
      assert torch.allclose(
          softmax(src, index, dim=1), src.softmax(dim=1)
      )
      if CALCULATION_VIA_PTR_AVAILABLE:
          assert torch.allclose(
              softmax(src, ptr=ptr, dim=1), src.softmax(dim=1)
          )
  
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 3 for calling a global function with example usage: test_softmax_dim() without return types