code: def learnable_p_choose(
      energy,
      noise_mean: float = 0.0,
      noise_var: float = 0.0,
      training: bool = True,
  ):
      """
      Calculating step wise prob for reading and writing
      1 to read, 0 to write
      energy: bsz, tgt_len, src_len
      """
  
      noise = 0
      if training:
          # add noise here to encourage discretness
          noise = (
              torch.normal(noise_mean, noise_var, energy.size())
              .type_as(energy)
              .to(energy.device)
          )
  
      p_choose = torch.sigmoid(energy + noise)
  
      # p_choose: bsz * self.num_heads, tgt_len, src_len
      return p_choose
  
imports: from typing import Optional, Dict
  from torch import Tensor
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: learnable_p_choose(energy, noise_mean, noise_var, training) and returns: p_choose