Method forward

axlearn/common/convolution.py:1807–1844

Stacks stride number of frames into one frame along the time axis. Args: inputs: Tensor of shape [batch, time, input_dim]. paddings: 0/1 boolean Tensor of shape [batch, time], paddings of the input sequences. Returns: stacked_inputs: Tensor of sh

(self, inputs: Tensor, *, paddings: Tensor)

        padding: Union[tuple[int, int], Literal["SAME", "VALID", "CAUSAL"]] = "VALID"

    def forward(self, inputs: Tensor, *, paddings: Tensor) -> tuple[Tensor, Tensor]:
        """Stacks stride number of frames into one frame along the time axis.

        Args:
            inputs: Tensor of shape [batch, time, input_dim].
            paddings: 0/1 boolean Tensor of shape [batch, time], paddings of the input sequences.

        Returns:
            stacked_inputs: Tensor of shape [batch, time // stride, input_dim * stride].
            stacked_paddings: 0/1 boolean Tensor of shape [batch, time // stride]. An output frame
                is padding if at least one of the stacked input frames is padding.

        Raises:
            ValueError: If stride is <= 1.
        """
        cfg = self.config
        if cfg.stride <= 1:
            raise ValueError(f"stride should be greater than 1, but got {cfg.stride}.")

        # For the last partial frame.
        inputs = inputs * safe_not(paddings)[:, :, None]

        padding = cfg.padding
        if isinstance(padding, str):
            padding = conv_explicit_padding(
                window=(cfg.stride,), strides=(cfg.stride,), padding=padding, dilation=(1,)
            )[0]
        inputs = jnp.pad(inputs, ((0, 0), padding, (0, 0)), constant_values=0)

        batch_size, seq_len, input_dim = inputs.shape
        output_length = seq_len // cfg.stride
        new_shape = [batch_size, output_length, input_dim * cfg.stride]
        # Stack inputs over the time dimension.
        stacked_inputs = jnp.reshape(inputs[:, : output_length * cfg.stride, :], new_shape)
        # An output frame is padding if at least one of the stacked input frames is padding.
        stacked_paddings = self.conv_paddings(paddings)
        stacked_inputs = stacked_inputs * safe_not(stacked_paddings)[:, :, None]
        return stacked_inputs, stacked_paddings

    @nowrap
    def conv_paddings(self, paddings: Tensor) -> Tensor: