code: def softmax(
      src: torch.Tensor,
      index: torch.LongTensor,
      num_nodes: Union[None, int, torch.Tensor] = None,
      dim: int = 0,
  ) -> torch.Tensor:
      r"""
      Compute a sparsely evaluated softmax.
  
      Given a value tensor :attr:`src`, this function first groups the values
      along the given dimension based on the indices specified in :attr:`index`,
      and then proceeds to compute the softmax individually for each group.
  
      :param src:
          The source tensor.
      :param index:
          The indices of elements for applying the softmax.
      :param num_nodes:
          The number of nodes, i.e., :obj:`max_val + 1` of :attr:`index`. (default: :obj:`None`)
      :param dim:
          The dimension along which to compute the softmax.
  
      :returns:
          The softmax-ed tensor.
  
      :raises ImportError: if :mod:`torch_scatter` is not installed
      """
      if torch_scatter is None:
          raise ImportError(
              "torch-scatter is not installed, attention aggregation won't work. "
              "Install it here: https://github.com/rusty1s/pytorch_scatter",
          )
      num_nodes = num_nodes or index.max() + 1
      out = src.transpose(dim, 0)
      out = (
          out
          - torch_scatter.scatter_max(
              out, index, dim=0, dim_size=num_nodes
          )[0][index]
      )
      out = out.exp()
      out = out / torch_scatter.scatter_add(
          out, index, dim=0, dim_size=num_nodes
      )[index].clamp_min(1.0e-16)
      return out.transpose(0, dim)
  
imports: from abc import abstractmethod
  from typing import ClassVar, Optional, Union
  import torch
  from class_resolver import ClassResolver
  from torch import nn
  from ..utils import einsum
label_desc: Use image node_id 1 for calling a global function with example usage: softmax(src, index, num_nodes, dim) and returns: out