flambe.nn.pooling

Module Contents

class flambe.nn.pooling.FirstPooling[source]

Bases: flambe.nn.Module

Get the last hidden state of a sequence.

forward(self, data: torch.Tensor, padding_mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The output data, as a tensor of shape [B x H]

Return type:

torch.Tensor

class flambe.nn.pooling.LastPooling[source]

Bases: flambe.nn.Module

Get the last hidden state of a sequence.

forward(self, data: torch.Tensor, padding_mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The output data, as a tensor of shape [B x H]

Return type:

torch.Tensor

class flambe.nn.pooling.SumPooling[source]

Bases: flambe.nn.Module

Get the sum of the hidden state of a sequence.

forward(self, data: torch.Tensor, padding_mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The output data, as a tensor of shape [B x H]

Return type:

torch.Tensor

class flambe.nn.pooling.AvgPooling[source]

Bases: flambe.nn.Module

Get the average of the hidden state of a sequence.

forward(self, data: torch.Tensor, padding_mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The output data, as a tensor of shape [B x H]

Return type:

torch.Tensor

flambe.nn.pooling._default_padding_mask(data: torch.Tensor) → torch.Tensor[source]

Builds a 1s padding mask taking into account initial 2 dimensions of input data.

Parameters:data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
Returns:A padding mask , as a tensor of shape [B x S]
Return type:torch.Tensor
flambe.nn.pooling._sum_with_padding_mask(data: torch.Tensor, padding_mask: torch.Tensor) → torch.Tensor[source]

Applies padding_mask and performs summation over the data

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The result of the summation, as a tensor of shape [B x H]

Return type:

torch.Tensor