docs: def _mask(self, tensor, length, padding_value=0):
      """Set padding elements of a batch of sequences to a constant.
  
      Useful for setting padding elements to zero before summing along the time
      dimension, or for preventing infinite results in padding elements.
  
      Args:
        tensor: Tensor of sequences.
        length: Batch of sequence lengths.
        padding_value: Value to write into padding elements.
  
      Returns:
        Masked sequences.
      """
      with tf.name_scope('mask'):
        range_ = tf.range(tensor.shape[1].value)
        mask = range_[None, :] < length[:, None]
        if tensor.shape.ndims > 2:
          for _ in range(tensor.shape.ndims - 2):
            mask = mask[..., None]
          mask = tf.tile(mask, [1, 1] + tensor.shape[2:].as_list())
        masked = tf.where(mask, tensor, padding_value * tf.ones_like(tensor))
        return tf.check_numerics(masked, 'masked')