Explicit scale management

Example source: examples/04_eager_scale_management.py

This example plans two plaintext-multiplication scales against the actual Q-group divisor, tracks every per-value scale transition, and aligns a depth separately. The tutorial also applies a bounded scale-metadata reinterpretation. Scale and depth transition laws are defined in Scale and depth lifecycle.

Run the example

Start with the shortest preset:

python examples/04_eager_scale_management.py --preset slots8192-scale40-depth7-int64

Track the actual scale

When no scale is supplied, encoding and encryption use:

Delta = engine.config.default_scale

Every value then owns its actual binary64 scale. Two plaintext products update it without consuming a depth:

$$\mathrm{\Delta }(c_{\mathrm{pre}})=\mathrm{\Delta }(c_{\mathrm{in}})\mathrm{\Delta }(p_1)\mathrm{\Delta }(p_2).$$

Public rescale_to_next_depth divides by the actual leading Q-group product:

$$\mathrm{\Delta }(c_{\mathrm{out}})=\frac{\mathrm{\Delta }(c_{\mathrm{pre}})}{M_{\mathrm{drop}}},M_{\mathrm{drop}}=\prod _{q\in G_d}q.$$

The result stores this quotient as its actual scale.

Query the arithmetic before encoding

The engine provides two pure transition queries:

drop_divisor = engine.rescale_divisor(depth=ciphertext.depth)
output_scale = engine.rescale_output_scale(
    input_scale=pre_rescale_scale,
    depth=ciphertext.depth,
)

rescale_divisor supplies the complete group product used to choose operand scales. rescale_output_scale calculates the actual output scale after the products have been evaluated.

For two plaintext products followed by one rescale_to_next_depth, targeting Delta gives:

first_scale = 2**20
second_scale = Delta * drop_divisor / (ciphertext.scale * first_scale)

The example uses this equation before encoding either operation-ready plaintext. Its planned branch reaches Delta after rescale_to_next_depth divides by the actual leading group product.

Align depth independently from scale

The original branch remains at depth zero. The planned product is at depth one. A modulus switch aligns the original value without changing its scale:

level_aligned = engine.mod_switch_to_depth(
    ciphertext,
    planned_product.depth,
)
combined = engine.add(planned_product, level_aligned)

This addition succeeds because the program independently arranged:

1. the same depth, through modulus switch;
2. the same scale, through plaintext-scale planning.

The operands therefore satisfy all addition preconditions before add is called.

Apply a bounded metadata reinterpretation

If the plaintext-scale product is Delta rather than the actual leading-group product $M_{\mathrm{drop}}$, rescale_to_next_depth reports:

$${\mathrm{\Delta }}^2/M_{\mathrm{drop}},$$

which differs from Delta by the actual-group-divisor ratio. A bounded reinterpretation records an application-approved target scale:

reinterpreted = engine.reinterpret_at_scale(
    actual_result,
    Delta,
    max_relative_change=1e-2,
)

The ciphertext residues remain unchanged. The decoded message is multiplied by old_scale / Delta, and the configured bound limits the accepted scale-ratio bias.

Source

#!/usr/bin/env python3

"""Plan and track actual per-value CKKS scales across two plaintext products.

Run:
    python examples/04_eager_scale_management.py --preset slots8192-scale40-depth7-int64
"""

from __future__ import annotations

import argparse
import math

import torch
from common import (
    add_engine_args,
    error_stats,
    make_engine,
    print_table,
    sync_if_cuda,
)

import fhelium as fh
from fhelium.eager import Engine


def multiply_twice_then_rescale_to_next_depth(
    engine: Engine,
    source: fh.Ciphertext,
    first_message: torch.Tensor,
    second_message: torch.Tensor,
    *,
    first_scale: float,
    second_scale: float,
) -> tuple[fh.Ciphertext, fh.Ciphertext, fh.Ciphertext]:
    r"""Apply two planned plaintext products followed by rescale-to-next.

    The actual-scale path is

    $$
    \Delta(c_1)=\Delta(c_0)\Delta(p_1),\qquad
    \Delta(c_2)=\Delta(c_1)\Delta(p_2),\qquad
    \Delta(c_3)=\frac{\Delta(c_2)}{M_{\mathrm{drop}}},
    \qquad M_{\mathrm{drop}}=\prod_{q\in G_d}q.
    $$

    Decoded slots satisfy $m_3\mathrel{\approx}m_0m_1m_2$ up to CKKS
    approximation error. Inputs are unchanged; all three returned ciphertexts
    own independent storage.

    Args:
        engine: Rank-local CKKS engine.
        source: Two-component coefficient-domain ciphertext.
        first_message: First cleartext slot-wise factor.
        second_message: Second cleartext slot-wise factor.
        first_scale: Encoding scale allocated to the first factor.
        second_scale: Encoding scale allocated to the second factor.

    Returns:
        Ciphertexts after the first product, second product, and rescale.
    """

    first = engine.prepare_plaintext_for_multiplication(
        engine.encode(first_message, depth=source.depth, scale=first_scale)
    )
    second = engine.prepare_plaintext_for_multiplication(
        engine.encode(second_message, depth=source.depth, scale=second_scale)
    )
    after_first = engine.multiply_plaintext(
        engine.coefficient_domain_to_ntt_domain(source), first
    )
    before_rescale = engine.multiply_plaintext(after_first, second)
    return (
        after_first,
        before_rescale,
        engine.rescale_to_next_depth(
            engine.ntt_domain_to_coefficient_domain(before_rescale)
        ),
    )


def scale_row(label: str, value: fh.Ciphertext) -> list[object]:
    """Return one display row for a ciphertext's recorded scale state."""

    return [
        label,
        value.depth,
        f"{value.scale:.17g}",
        f"{math.log2(value.scale):.9f}",
    ]


def error_row(
    label: str,
    engine: Engine,
    value: fh.Ciphertext,
    expected: torch.Tensor,
) -> list[object]:
    """Decrypt one result and return compact approximation-error statistics."""

    error = error_stats(
        engine.decrypt_message(value).cpu(),
        expected,
        engine.num_slots,
    )
    return [label, f"{error['max_abs']:.3e}", f"{error['rms']:.3e}"]


def main() -> None:
    """Run the scale-planning and guarded-reinterpretation example."""

    parser = argparse.ArgumentParser(description=__doc__)
    add_engine_args(parser, default_preset="slots8192-scale40-depth7-int64")
    parser.add_argument(
        "--reinterpret-bound",
        type=float,
        default=1e-2,
        help=(
            "Maximum symmetric relative scale change accepted by the guarded "
            "reinterpretation."
        ),
    )
    args = parser.parse_args()

    engine = make_engine(args)
    slots = engine.num_slots
    default_scale = engine.config.default_scale

    message = 0.01 * torch.sin(torch.arange(slots, dtype=torch.float64) * 0.013)
    first_message = torch.linspace(0.8, 1.2, slots, dtype=torch.float64)
    second_message = torch.linspace(1.1, 0.9, slots, dtype=torch.float64)
    expected_product = message * first_message * second_message

    source = engine.encrypt_message(message, scale=default_scale)
    dropped_divisor = engine.rescale_divisor(depth=source.depth)

    # The first factor receives a precision allocation. The second
    # scale completes the product required to reach default_scale after
    # division by the actual leading Q-group product.
    default_scale_bits = math.floor(math.log2(default_scale))
    first_scale = float(1 << (default_scale_bits // 2))
    second_scale = default_scale * dropped_divisor / (source.scale * first_scale)
    planned_first, planned_pre, planned = (
        multiply_twice_then_rescale_to_next_depth(
            engine,
            source,
            first_message,
            second_message,
            first_scale=first_scale,
            second_scale=second_scale,
        )
    )
    predicted_scale = engine.rescale_output_scale(
        planned_pre.scale,
        depth=planned_pre.depth,
    )
    assert planned.scale == predicted_scale == default_scale

    # Depth alignment is independent of scale alignment. Because the planned
    # branch reaches default_scale exactly, a depth-only modulus switch makes
    # the original branch directly add-compatible.
    depth_aligned_source = engine.mod_switch_to_depth(source, planned.depth)
    planned_sum = engine.add(planned, depth_aligned_source)

    # The comparison allocation uses a plaintext-scale product of Delta.
    # Rescale records Delta^2 divided by the leading Q-group product, so the
    # result does not meet the matching-scale precondition for addition.
    approximate_second_scale = default_scale / first_scale
    approximate_first, approximate_pre, approximate = (
        multiply_twice_then_rescale_to_next_depth(
            engine,
            source,
            first_message,
            second_message,
            first_scale=first_scale,
            second_scale=approximate_second_scale,
        )
    )
    scale_difference = approximate.scale / depth_aligned_source.scale - 1.0

    # Guarded reinterpretation records the configured target scale while
    # preserving residues. The bound limits the accepted message bias.
    reinterpreted = engine.reinterpret_at_scale(
        approximate,
        default_scale,
        max_relative_change=args.reinterpret_bound,
    )
    reinterpreted_sum = engine.add(reinterpreted, depth_aligned_source)
    sync_if_cuda(torch.get_default_device())

    print(engine)
    print(f"default scale Delta: {default_scale:.17g}")
    print(f"depth-0 rescale divisor: {dropped_divisor}")
    print(
        f"allocated plaintext scales: p1={first_scale:.17g}, p2={second_scale:.17g}"
    )
    print("\nExplicit scale states")
    print_table(
        ["value", "depth", "scale", "log2(scale)"],
        [
            scale_row("source", source),
            scale_row("planned after pmult 1", planned_first),
            scale_row("planned before rescale", planned_pre),
            scale_row("planned after rescale", planned),
            scale_row("depth-only switched source", depth_aligned_source),
            scale_row("Delta-product before rescale", approximate_pre),
            scale_row("actual Delta^2/M0 result", approximate),
            scale_row("explicitly reinterpreted", reinterpreted),
        ],
    )

    print("\nScale comparison before explicit reinterpretation")
    print(f"Relative scale difference: {scale_difference:.9e}")
    print("The caller aligns actual scales before addition; Eager does not infer alignment.")

    # Reinterpretation changes the decoded product by old_scale/new_scale.
    reinterpret_ratio = approximate.scale / reinterpreted.scale
    print("\nCleartext error")
    print_table(
        ["result", "max abs error", "rms error"],
        [
            error_row("planned product", engine, planned, expected_product),
            error_row(
                "actual-scale Delta-product",
                engine,
                approximate,
                expected_product,
            ),
            error_row(
                "planned product + depth-switched source",
                engine,
                planned_sum,
                expected_product + message,
            ),
            error_row(
                "reinterpreted product + source",
                engine,
                reinterpreted_sum,
                expected_product * reinterpret_ratio + message,
            ),
        ],
    )


if __name__ == "__main__":
    main()

Next step

Continue with Late relinearization and NTT reuse to combine scale transitions with three-component products and representation reuse. The complete transition laws are defined in Scale and depth lifecycle, and the parameter invariants are described in Configuration and modulus chain.