def encode(coords: torch.Tensor, permute: List[int] = [0, 1, 2], mode: Literal['z_order', 'hilbert'] = 'z_order') -> torch.Tensor:
    """
    Encodes 3D coordinates into a 30-bit code.

    Args:
        coords: a tensor of shape [N, 3] containing the 3D coordinates.
        permute: the permutation of the coordinates.
        mode: the encoding mode to use.
    """
    if mode == 'z_order':
        return z_order_encode(coords[:, permute], depth=10).int()
    elif mode == 'hilbert':
        return hilbert_encode(coords[:, permute], depth=10).int()
    else:
        raise ValueError(f"Unknown encoding mode: {mode}")

    Encodes 3D coordinates into a 30-bit code.

    Args:
        coords: a tensor of shape [N, 3] containing the 3D coordinates.
        permute: the permutation of the coordinates.
        mode: the encoding mode to use.
    
extensions/vox2seq/vox2seq/pytorch/__init__.py