code: def softmax(
      src: Tensor,
      index: Optional[Tensor] = None,
      ptr: Optional[Tensor] = None,
      num_nodes: Optional[int] = None,
      dim: int = 0,
  ) -> Tensor:
      if ptr is not None:
          dim = dim + src.dim() if dim < 0 else dim
          size = ([1] * dim) + [-1]
          ptr = ptr.view(size)
          src_max = gather_csr(segment_csr(src, ptr, reduce="max"), ptr)
          out = (src - src_max).exp()
          out_sum = gather_csr(segment_csr(out, ptr, reduce="sum"), ptr)
      elif index is not None:
          N = maybe_num_nodes(index, num_nodes)
          src_max = scatter(src, index, dim, dim_size=N, reduce="max")
          src_max = src_max.index_select(dim, index)
          out = (src - src_max).exp()
          out_sum = scatter(out, index, dim, dim_size=N, reduce="sum")
          out_sum = out_sum.index_select(dim, index)
      else:
          raise NotImplementedError
  
      return out / (out_sum + 1e-16)
  
imports: import torch
  from torch import nn
  import torch.nn.functional
  import dgl
  from dgl.nn.pytorch.conv import *
  import dgl.function
  import math
  from dgl.utils import expand_as_pair
  from dgl.nn.functional import edge_softmax
  from torch.nn import Parameter
  import torch
  from torch import Tensor
  from ..scatter_utils import *
  from typing import Optional, Tuple, Union
label_desc: Use image node_id 4 for calling a global function with example usage: softmax(src, index, ptr, num_nodes, dim) and returns: unknown