import torch


def encode_dist(x, scales=[1,2,4,8], include_self = True):
    """ Encodes a distance with sines and cosines. 
        Inputs:
        * x: (batch, N) or (N,). data to encode.
              Infer devic and type (f16, f32, f64) from here.
        * scales: (s,) or list. lower or higher depending on distances.
        Output: (..., num_scales*2 + 1) if include_self or (..., num_scales*2) 
    """
    x = x.unsqueeze(-1)
    # infer device
    device, precise = x.device, x.type()
    # convert to tensor
    if isinstance(scales, list):
        scales = torch.tensor([scales], device=device).type(precise)
    # get pos encodings
    sines   = torch.sin(x / scales)
    cosines = torch.cos(x / scales)
    # concat and return
    enc_x = torch.cat([sines, cosines], dim=-1)
    return torch.cat([enc_x, x], dim=-1) if include_self else enc_x