code: def masked_softmax(
      vector: torch.Tensor, mask: torch.Tensor, dim: int = -1
  ) -> torch.Tensor:
      """
      ``torch.nn.functional.softmax(vector)`` does not work if some elements of ``vector`` should be
      masked.  This performs a softmax on just the non-masked portions of ``vector``.  Passing
      ``None`` in for the mask is also acceptable; you'll just get a regular softmax.
  
      ``vector`` can have an arbitrary number of dimensions; the only requirement is that ``mask`` is
      broadcastable to ``vector's`` shape.  If ``mask`` has fewer dimensions than ``vector``, we will
      unsqueeze on dimension 1 until they match.  If you need a different unsqueezing of your mask,
      do it yourself before passing the mask into this function.
  
      In the case that the input vector is completely masked, this function returns an array
      of ``0.0``. This behavior may cause ``NaN`` if this is used as the last layer of a model
      that uses categorical cross-entropy loss.
      """
      if mask is None:
          result = torch.nn.functional.softmax(vector, dim=dim)
      else:
          mask = mask.float()
          while mask.dim() < vector.dim():
              mask = mask.unsqueeze(1)
          # To limit numerical errors from large vector elements outside the mask, we zero these out.
          result = torch.nn.functional.softmax(vector * mask, dim=dim)
          result = result * mask
          result = result / (result.sum(dim=dim, keepdim=True) + 1e-13)
      return result
  
file_path: jiant/jiant/ext/allennlp.py
imports: from typing import Optional
  import torch
  import torch.nn
label_desc: Use image node_id 2 for calling a global function with example usage: masked_softmax(vector, mask, dim) and returns: result