code: def shuffle_node(
      x: Tensor,
      batch: Optional[Tensor] = None,
      training: bool = True,
  ) -> Tuple[Tensor, Tensor]:
      r"""Randomly shuffle the feature matrix :obj:`x` along the
      first dimmension.
  
      The method returns (1) the shuffled :obj:`x`, (2) the permutation
      indicating the orders of original nodes after shuffling.
  
      Args:
          x (FloatTensor): The feature matrix.
          batch (LongTensor, optional): Batch vector
              :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
              node to a specific example. Must be ordered. (default: :obj:`None`)
          training (bool, optional): If set to :obj:`False`, this operation is a
              no-op. (default: :obj:`True`)
  
      :rtype: (:class:`FloatTensor`, :class:`LongTensor`)
  
      Example:
          >>> # Standard case
          >>> x = torch.tensor([[0, 1, 2],
          ...                   [3, 4, 5],
          ...                   [6, 7, 8],
          ...                   [9, 10, 11]], dtype=torch.float)
          >>> x, node_perm = shuffle_node(x)
          >>> x
          tensor([[ 3.,  4.,  5.],
                  [ 9., 10., 11.],
                  [ 0.,  1.,  2.],
                  [ 6.,  7.,  8.]])
          >>> node_perm
          tensor([1, 3, 0, 2])
  
          >>> # For batched graphs as inputs
          >>> batch = torch.tensor([0, 0, 1, 1])
          >>> x, node_perm = shuffle_node(x, batch)
          >>> x
          tensor([[ 3.,  4.,  5.],
                  [ 0.,  1.,  2.],
                  [ 9., 10., 11.],
                  [ 6.,  7.,  8.]])
          >>> node_perm
          tensor([1, 0, 3, 2])
      """
      if not training:
          perm = torch.arange(x.size(0), device=x.device)
          return x, perm
      if batch is None:
          perm = torch.randperm(x.size(0), device=x.device)
          return x[perm], perm
      num_nodes = scatter(
          batch.new_ones(x.size(0)), batch, dim=0, reduce="sum"
      )
      ptr = cumsum(num_nodes)
      perm = torch.cat(
          [
              torch.randperm(n, device=x.device) + offset
              for offset, n in zip(ptr[:-1], num_nodes)
          ]
      )
      return x[perm], perm
  
imports: from typing import Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch_geometric.utils import cumsum, negative_sampling, scatter
label_desc: Use image node_id 1 for calling a global function with example usage: shuffle_node(x, batch, training) and returns: x,  perm, x,  perm, x,  perm