torchwrench.nn.functional.mask module

torchwrench.nn.functional.mask.generate_square_subsequent_mask(size: int, diagonal: int = 0, *, device: device | None | 'default' | 'cuda_if_available' | str | int = None, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]
torchwrench.nn.functional.mask.lengths_to_non_pad_mask(lengths: Tensor, max_len: int | None = None, include_end: bool = False, *, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]

Convert lengths to binary mask of non-padded values.

The output will be a tensor of shape (B, max_len).

Args:

lengths: (bsize,) max_len: Optional int for indicate the maximal length.

If None, it will be set to lengths.max(). defaults to None.

include_end: If True, the value at index of len will be True in returned mask.

defaults to False.

Example 1::

>>> input = torch.as_tensor([4, 2, 0, 3, 0])
>>> lengths_to_non_pad_mask(input, max_len=6, include_end=False)
tensor([[True, True, True, True, False, False],
        [True, True, False, False, False, False],
        [False, False, False, False, False, False],
        [True, True, True, False, False, False],
        [False, False, False, False, False, False]])
torchwrench.nn.functional.mask.lengths_to_pad_mask(lengths: Tensor, max_len: int | None = None, include_end: bool = True, *, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]

Convert lengths to binary mask of padded values. The output will be a tensor of shape (B, max_len).

Args:

lengths: (B,) max_len: Optional int for indicate the maximal length.

If None, it will be set to lengths.max(). defaults to None.

include_end: If True, the last value of each size will be set to False.

defaults to True.

Example 1::

>>> input = torch.as_tensor([4, 2, 0, 3, 0])
>>> lengths_to_non_pad_mask(input, max_len=None, include_end=True)
tensor([[False, False, False, False],
        [False, False, True, True],
        [True, True, True, True],
        [False, False, False, True],
        [True, True, True, True]])
torchwrench.nn.functional.mask.lengths_to_ratios(lengths: Tensor, max_len: int | None = None) Tensor[source]
torchwrench.nn.functional.mask.masked_equal(x1: Tensor, x2: Tensor, mask: Tensor) bool[source]

Check if two tensors are equal at the specific positions.

Args:

x1: First tensor of shape S. x2: Second tensor of shape S. mask: Boolean tensor of shape S. Position marked as False are ignored by the equality.

torchwrench.nn.functional.mask.masked_mean(x: T_TensorOrArray, non_pad_mask: T_TensorOrArray, *, dim: None | int | Iterable[int] = None, min_div: float | None = 1.0) T_TensorOrArray[source]

Average a tensor along the specified dim(s).

Args:

tensor: (N, …) non_pad_mask: Non-padding mask, should be broadcastable with argument tensor and reduced with argument dim.

It should be a boolean tensor or a float tensor containing only 1 and 0 values.

dim: Optional dim(s) to reduce. If None, result will be reduced to a scalar. defaults to None. min_div: Minimal value to avoid division by 0. defaults to 1.0.

torchwrench.nn.functional.mask.masked_sum(x: T_TensorOrArray, non_pad_mask: T_TensorOrArray, *, dim: None | int | Iterable[int] = None) T_TensorOrArray[source]

Sum a tensor along the specified dim(s).

Args:

x: (N, …) non_pad_mask: Non-padding mask, should be broadcastable with argument tensor and reduced with argument dim.

It should be a boolean tensor or a float tensor containing only 1 and 0 values.

dim: Optional dim(s) to reduce. If None, result will be reduced to a scalar. defaults to None.

torchwrench.nn.functional.mask.non_pad_mask_to_lengths(mask: T_TensorOrArray, *, dim: int = -1) T_TensorOrArray[source]
torchwrench.nn.functional.mask.non_pad_mask_to_ratios(non_pad_mask: Tensor, *, dim: int = -1) Tensor[source]
torchwrench.nn.functional.mask.pad_mask_to_lengths(mask: T_TensorOrArray, *, dim: int = -1) T_TensorOrArray[source]
torchwrench.nn.functional.mask.pad_mask_to_ratios(pad_mask: Tensor, *, dim: int = -1) Tensor[source]
torchwrench.nn.functional.mask.ratios_to_lengths(ratios: Tensor, max_len: int, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]
torchwrench.nn.functional.mask.ratios_to_non_pad_mask(ratios: Tensor, max_len: int, include_end: bool = False, *, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]
torchwrench.nn.functional.mask.ratios_to_pad_mask(ratios: Tensor, max_len: int, include_end: bool = True, *, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]
torchwrench.nn.functional.mask.tensor_to_lengths(tensor: Tensor, *, pad_value: float | None = None, end_value: float | None = None, dim: int = -1) LongTensor[source]

Get the lengths of the non-padded elements of a tensor.

You must provide a value for one of pad_value or end_value. If both values are provided, the end_value is ignored. The output will be of shape (N,). The end_value is not included in the length of the sentence.

Args:

tensor: Input of shape (N, *). pad_value: The pad value used in tensor. defaults to None. end_value: The end value used in tensor. defaults to None. dim: The dimension of the length. defaults to -1.

Example 1::

` >>> x = torch.as_tensor([1, 10, 20, 2, 0, 0]) >>> tensor_to_lengths(x, end_value=2) ... tensor(3) `

Example 2::

` >>> x = torch.as_tensor([1, 10, 20, 2, 0, 0]) >>> tensor_to_lengths(x, pad_value=0) ... tensor(4) `

torchwrench.nn.functional.mask.tensor_to_non_pad_mask(tensor: Tensor, *, pad_value: float | None = None, end_value: float | None = None, include_end: bool = False, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]

Convert tensor to non-pad binary mask. You must provide a value for one of pad_value or end_value. If both values are provided, the end_value is ignored. The output will be a binary mask representing the non-padded values. Shape is the same than the input tensor.

Args:

tensor: A tensor of values. If end_value is given instead of pad_value, the number of dims must be <= 2. pad_value: The pad value used in tensor. defaults to None. end_value: The end value used in tensor. defaults to None. include_end: If True, the end value will be included in non_pad_mask.

This parameter is ignored if end_value is None. defaults to False.

Example 1::

>>> input = torch.as_tensor([1, 10, 20, 2, 0, 0])
>>> tensor_to_pad_mask(input, end_value=2)
tensor([True, True, True, False, False, False])
torchwrench.nn.functional.mask.tensor_to_pad_mask(tensor: Tensor, *, pad_value: float | None = None, end_value: float | None = None, include_end: bool = True, dtype: dtype | None | 'default' | str | DTypeEnum = None) Tensor[source]

Convert tensor to pad binary mask.

You must provide a value for one of pad_value or end_value. If both values are provided, the end_value is ignored.

The output will be a binary mask representing the padded values. Shape is the same than the input tensor.

Args:

tensor: A tensor of values. If end_value is given instead of pad_value, the number of dims must be <= 2. pad_value: The pad value used in tensor. defaults to None. end_value: The end value used in tensor. defaults to None. include_end: If True, the end value will be included in pad_mask. defaults to True.

Example 1::

>>> input = torch.as_tensor([1, 10, 20, 2, 0, 0])
>>> tensor_to_pad_mask(input, end_value=2)
tensor([False, False, False, True, True, True])
torchwrench.nn.functional.mask.tensor_to_tensors_list(x: Tensor, *, pad_value: float | None = None, end_value: float | None = None, non_pad_mask: Tensor | None = None, lengths: None | Tensor | list[int] = None, dim: int = -1) list[Tensor][source]

Convert padded tensor to tensor list.

You must provide a value for one of the 4 arguments: pad_value, end_value, non_pad_mask or lengths. If multiple values are provided, only one will be used and the priority order is pad_value, end_value, non_pad_mask and lengths. The output will be a list of N tensors of shape (*).

Args:

tensor: (N, *) pad_value: Pad value index. defaults to None. end_value: End value index. defaults to None. non_pad_mask: Optional non-padded boolean mask. defaults to None. lengths: Length of each sequence in padded batch. dim: Dimension to get lengths. defaults to -1.

torchwrench.nn.functional.mask.tensors_list_to_lengths(tensors: list[Tensor], dim: int = -1) LongTensor1D[source]

Return the size of the tensor at a specific dim.

The output will be a tensor of size N.

Args:

tensors: List of N tensors. dim: The dimension of the output sizes. defaults to -1.