Simple NeRF positional embedding

positional_embedding.py

# simple positional embedding (NeRF style)
# as in https://arxiv.org/pdf/2003.08934 (Equation 4)

def positional_embedding(x, k):

    batch_size, n_dims = x.shape
    x_rep = x.unsqueeze(-1).expand(-1, -1, k)

    k_pows = torch.pow(2, torch.arange(0, k, device=x.device)).view(1, 1, k)
    x_k = k_pows * torch.pi * x_rep

    x_pos = torch.cat([torch.sin(x_k), torch.cos(x_k)], dim=-1)

    return x_pos.view(batch_size, -1)