fhelium.engine.slot_embedding

Low-level CKKS slot/polynomial encoding helpers.

make_slot_tensor

function View source

def make_slot_tensor(m: Sequence | torch.Tensor | complex, num_slots: int, device: str | torch.device='cpu') -> torch.Tensor: ...

Materialize the canonical [*batch, slot] message layout.

A scalar is repeated across all $S$ slots. A tensor-like message with final extent at most $S$ preserves every leading batch axis and is padded with zeros on the right of the slot axis. Existing tensor dtype is preserved; Python/non-tensor inputs are inferred by PyTorch. device places scalar and non-tensor inputs, while an existing non-scalar tensor retains its current device until the embedding operation moves it.

The function is functional for scalar/non-tensor inputs and for inputs requiring padding. A full-length tensor may be returned directly and can therefore alias the input.

encode_slots

function View source

def encode_slots(m: torch.Tensor, rng: Csprng, scale: float=float(2 ** 40), device: str | torch.device='cpu', galois_generator: int=3) -> torch.Tensor: ...

Encode canonical slots as stochastic-rounded integer coefficients.

For $m\in {\mathbb{C}}^S$ in the slot order selected by galois_generator, compute

$$p_i=\mathrm{SRound}(\mathrm{\Delta }{\mathcal{E}}_g^{-1}(m{)}_i),\mathbb{E}[\mathrm{SRound}(x)]=x.$$

m has layout [*batch, slot] with final extent $S=N/2$. The result has layout [*batch, coefficient] with final extent $N$, uses the integral dtype returned by rng.randround, and resides on device. It is a coefficient-domain polynomial, not RNS data. The operation is functional and the output does not alias m.

inverse_embed_slots

function View source

def inverse_embed_slots(m: torch.Tensor, *, device: str | torch.device='cpu', galois_generator: int=3) -> torch.Tensor: ...

Apply the unscaled inverse canonical embedding ${\mathcal{E}}_g^{-1}$.

m is [*batch, slot] in FHElium's canonical generator-3 order or the conventional OpenFHE generator-5 order, with $S=N/2$. The implementation permutes slots into a conjugate-symmetric length-$N$ vector $z$, then uses

$$a_k=\mathrm{Re}(e^{-\pi ik/N}{\mathrm{FFT}}_{\mathrm{forward}}(z{)}_k).$$

PyTorch norm="forward" contributes the $1/N$ forward normalization; this is paired with the inverse transform in embed_coefficients. The result is a binary64 real [*batch, coefficient] tensor on device. Leading batch axes are preserved. It has not been scaled, quantized, or reduced into RNS, and it does not alias m.

decode_slots

function View source

def decode_slots(m: torch.Tensor, scale: float=float(2 ** 40), galois_generator: int=3) -> torch.Tensor: ...

Decode coefficient values with their actual CKKS scale.

For m with layout [*batch, coefficient] and final extent $N$,

$$m_{\mathrm{slots}}={\mathcal{E}}_g(m)/\mathrm{\Delta }.$$

The result is complex binary64 [*batch, canonical_slot] with final extent $N$ on the input device; callers select the first $S=N/2$ semantic slots. m may contain exact integer coefficients or bounded binary64 approximate decrypt coefficients. The operation is functional.

embed_coefficients

function View source

def embed_coefficients(m: torch.Tensor, *, galois_generator: int=3) -> torch.Tensor: ...

Apply the unscaled canonical embedding ${\mathcal{E}}_g$.

For a real [*batch, coefficient] tensor $a$ with final extent $N$, form

$$z={\mathrm{IFFT}}_{\mathrm{forward}}\left(a_k{e}^{\pi ik/N}\right)$$

and permute $z$ into the selected canonical slot order. PyTorch norm="forward" makes the inverse transform unnormalized so this is the inverse of inverse_embed_slots. The result is complex binary64, preserves all leading batch axes, remains on the input device, and does not alias m.

circular_shift_permutation

function View source

def circular_shift_permutation(N: int, shift: int=1) -> torch.Tensor: ...

canon_permutation

function View source

def canon_permutation(N: int, k: int=1, verbose: bool=False) -> torch.Tensor: ...

Permutes the coefficients of the lattice basis that yields correctly the permutation of the decoded message.

The canonical permutation is defined as mu_p(n) = pn mod M where p is coprime with M, where p=2*k+1.

fold_permutation

function View source

def fold_permutation(N: int, p: torch.Tensor, verbose: bool=False) -> torch.Tensor: ...

In application to crypto, we fold the FFT at Nyquist.

Inverse FFT results in selection of alternating elements. Folding should correct the indices of the permutation according to the folding rule.

For example, 1->0, 3->1, 5->2, and so on.

conjugate_permutation

function View source

def conjugate_permutation(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor: ...

Conjugate permutations p and q by stacking p on top of q.

Permutations p and q must share the same cycle structures.

permutation_cycles

function View source

def permutation_cycles(perm: torch.Tensor) -> list[list[int]]: ...

Transform a plain permutation into a composition of cycles.

inverse_permutation

function View source

def inverse_permutation(p: torch.Tensor, verbose: bool=False) -> torch.Tensor: ...

Calculates the inverse permutation.

expand2conjugate

function View source

def expand2conjugate(m: torch.Tensor) -> torch.Tensor: ...

generate_twister

function View source

def generate_twister(N: int, device: str | torch.device='cpu') -> torch.Tensor: ...

generate_skewer

function View source

def generate_skewer(N: int, device: str | torch.device='cpu') -> torch.Tensor: ...

m2poly

function View source

def m2poly(m: torch.Tensor, twister: torch.Tensor) -> torch.Tensor: ...

m is the message and this function turns the message into polynomial coefficients. The message must be expanded mirrored in conjugacy.

poly2m

function View source

def poly2m(poly: torch.Tensor, skewer: torch.Tensor) -> torch.Tensor: ...

poly is the polynomial coefficients and this function turns the coefficients into a plain message.

prepost_perms

function View source

def prepost_perms(N: int, device: str | torch.device='cpu', galois_generator: int=3) -> tuple[torch.Tensor, torch.Tensor]: ...

pre_permute

function View source

def pre_permute(m: torch.Tensor, pre_perm: torch.Tensor) -> torch.Tensor: ...

Input m must be a torch tensor.

post_permute

function View source

def post_permute(m: torch.Tensor, post_perm: torch.Tensor) -> torch.Tensor: ...

Input m must be a torch tensor.