Function positional_encoding

pointersect/models/model_utils.py:199–250

r"""Apply positional encoding to the input. Args: tensor: (*, dim_in) Input tensor to be positionally encoded. num_encoding_functions: Number of encoding functions used to compute a positional encoding (default: 6). include_input: Whet

(
        tensor: torch.Tensor,
        num_encoding_functions: int = 6,
        include_input: bool = True,
        log_sampling: bool = True,
)

def positional_encoding(
        tensor: torch.Tensor,
        num_encoding_functions: int = 6,
        include_input: bool = True,
        log_sampling: bool = True,
) -> torch.Tensor:
    r"""Apply positional encoding to the input.

    Args:
        tensor:
            (*, dim_in) Input tensor to be positionally encoded.
        num_encoding_functions:
            Number of encoding functions used to compute a positional encoding (default: 6).
        include_input:
            Whether to include the input in the positional encoding (default: True).
        log_sampling:
            whether to sample the sinusoid frequencies in log scale.

    Returns:
    (torch.Tensor): Positional encoding of the input tensor.  (*, dim_out)

    .. math::
        dim_{out} = d_{in} * include_input + num_encoding_functions * 2
    """
    encoding = [tensor] if include_input else []
    frequency_bands = None
    if log_sampling:
        frequency_bands = 2.0 ** torch.linspace(
            0.0,
            num_encoding_functions - 1,
            num_encoding_functions,
            dtype=tensor.dtype,
            device=tensor.device,
        )
    else:
        frequency_bands = torch.linspace(
            2.0 ** 0.0,
            2.0 ** (num_encoding_functions - 1),
            num_encoding_functions,
            dtype=tensor.dtype,
            device=tensor.device,
        )

    for freq in frequency_bands:
        for func in [torch.sin, torch.cos]:
            encoding.append(func(tensor * freq))

    # Special case, for no positional encoding
    if len(encoding) == 1:
        return encoding[0]
    else:
        return torch.cat(encoding, dim=-1)


def get_embedding_function(