code: def gumbel_softmax(
      log_alpha: torch.Tensor, beta: float = 1.0, training: bool = True
  ):
      r"""Sample from the instantiation of concrete distribution when training
      Args:
          log_alpha: input probabilities
          beta: temperature for softmax
      """
      if training:
          random_noise = torch.rand(log_alpha.shape)
          random_noise = torch.log(random_noise) - torch.log(
              1.0 - random_noise
          )
          gate_inputs = (
              random_noise.to(log_alpha.device) + log_alpha
          ) / beta
          gate_inputs = gate_inputs.sigmoid()
      else:
          gate_inputs = log_alpha.sigmoid()
  
      return gate_inputs
  
file_path: DIG/dig/xgraph/models/utils.py
imports: import torch
  from torch_geometric.utils.num_nodes import maybe_num_nodes
label_desc: Use image node_id 3 for calling a global function with example usage: gumbel_softmax(log_alpha, beta, training) and returns: gate_inputs