Function reconstruct

mom/layers/mom.py:206–277

Reconstruct and mix transformed outputs back into the original input sequence shape. Key operations: 1. Reshapes and transposes `transformed_x` to prepare for scattering. 2. Applies the `mask` to zero out invalid positions. 3. Uses `torch.scatter_add_` to scatter and sum the tr

(
    transformed_x,
    indices: torch.Tensor,
    sorted_indices: torch.Tensor,
    batch_size: int,
    seq_len: int,
    topk: int,
    routing_weights: torch.Tensor,
    mask: torch.Tensor
)

def reconstruct(
    transformed_x,
    indices: torch.Tensor,
    sorted_indices: torch.Tensor,
    batch_size: int,
    seq_len: int,
    topk: int,
    routing_weights: torch.Tensor,
    mask: torch.Tensor
):
    '''
    Reconstruct and mix transformed outputs back into the original input sequence shape.

    Key operations:
    1. Reshapes and transposes `transformed_x` to prepare for scattering.
    2. Applies the `mask` to zero out invalid positions.
    3. Uses `torch.scatter_add_` to scatter and sum the transformed outputs back to their original positions based on `indices`.
    4. Rearranges the scattered outputs using `sorted_indices` to ensure correct ordering.
    5. Applies the `routing_weights` to weight the outputs.
    6. Sums over the `topk` dimension to produce the final reconstructed output.

    Args:
        transformed_x (torch.Tensor):
            The transformed output tensor from memory units or experts.
            Shape: (num_memories, batch_size, capacity_len, hidden_size)
        indices (torch.Tensor):
            Indices used for scattering the transformed outputs back to their corresponding positions.
            Shape: (batch*num_memories, max_len)
        sorted_indices (torch.Tensor):
            Sorting indices used to rearrange the scattered outputs back into the original sequence order.
            Shape: (batch_size*seq_len*topk)
        batch_size (int):
            The size of the batch.
        seq_len (int):
            The length of the input sequence.
        topk (int):
            The number of top elements selected (`topk`) per token during the selection process.
        routing_weights (torch.Tensor):
            Routing weights assigned to the top-k selected outputs when reconstructing the final output.
            Shape: (batch_size, seq_len, topk)
        mask (torch.Tensor):
            Boolean mask indicating valid positions in the sequence.
            Shape: (batch*num_memories, max_len)

    Returns:
        restored_x (torch.Tensor):
            The reconstructed output tensor in the original input sequence shape.
            Shape: (batch_size, seq_len, hidden_size)
    '''
    transformed_x = transformed_x.transpose(0, 1).reshape(
        (-1, transformed_x.shape[2], transformed_x.shape[3]))
    b, s, k, d = batch_size, seq_len, topk, transformed_x.shape[2]
    gathered_x = transformed_x.reshape(
        (transformed_x.shape[0] * transformed_x.shape[1], transformed_x.shape[2]))
    mask_expanded = mask.reshape(-1).unsqueeze(-1).expand_as(gathered_x)
    gathered_x = gathered_x * mask_expanded

    assert (indices >= 0).all(), "Indices should be non-negative"