Slots, rotations & stages
On Ciphertext state we saw that one axis of your tensor, the n axis, is packed into a ciphertext's slots. Values in different slots can't reach each other directly: any operator that needs to copy or sum values across slots must rotate the ciphertext, and rotations are among the most expensive operators a pipeline does. Stages, set through the stage_sizes parameter, are how you control that cost.
One axis lives in slots
HomParams.internal_n is the ciphertext's physical slot capacity (n / 2 in CKKS), where n is the ring dimension from Choosing parameters. One axis of your tensor, the n axis, goes into those slots. HomValue.n_slots is the block that axis is packed into; it defaults to internal_n. How packing fills those slots is on Ciphertext state.HomAdd between two ciphertexts. Along the n axis it is not: the values are sealed inside one encrypted object, and nothing can read "slot 5".Element-wise operators touch every slot independently and in parallel. The packing costs them nothing.
To sum the n axis, or copy a value along it, slots must line up with each other. The only way to do that is to shift the whole ciphertext, a rotation.
Rotations
slots [ x0 x1 x2 x3 x4 x5 x6 x7 ] rotate by 2 [ x2 x3 x4 x5 x6 x7 x0 x1 ] sum [ x0+x2 x1+x3 x2+x4 x3+x5 x4+x6 x5+x7 x6+x0 x7+x1 ]
Stages: stage_sizes
k rows, a column sum of n_slots / nb_cols slots, or an expand across input.n_slots. They do that work in rounds, and one round is a stage. You choose the stage plan as a plain list of integers, one entry per round: stage_sizes. Each entry is that stage's group size, how many pieces it combines at once. Summing 8 slots with stage_sizes = (2, 2, 2) means three rounds, each combining pairs:stage 1 (size 2) (x0+x1) (x2+x3) (x4+x5) (x6+x7) covers 2 slots each stage 2 (size 2) (x0+x1+x2+x3) (x4+x5+x6+x7) covers 4 slots each stage 3 (size 2) (x0+x1+x2+x3+x4+x5+x6+x7) covers 8 = k
HomRunningSum that is prod(stage_sizes) == k. For HomSumSlots it is prod(stage_sizes) == n_slots / nb_cols. For HomExpand it is prod(stage_sizes) == input.n_slots. Plans do not have to be uniform: (4, 2) also covers 8. (3, 2) does not: 3 × 2 = 6, so compilation fails.s needs s − 1 distinct rotation offsets. Summing those across every stage gives the plan's total offset count, the number of distinct rotation-key entries the plan needs overall:(2, 2, 2)1 + 1 + 1 = 3, over three rounds(4, 2)3 + 1 = 4, over two rounds(8,)7, in a single roundk = 1024, the plan (2,)×10 needs 10 offsets over ten rounds; (1024,) needs 1,023 in one round; (32, 32) sits in between with 62 over two.Defaults: leave stage_sizes unset and each round has size 2. That is (2,) × log₂(k) for HomRunningSum, (2,) × log₂(n_slots / nb_cols) for HomSumSlots, and (2,) × log₂(input.n_slots) for HomExpand. This plan only exists when the span is a power of two. Otherwise write the plan yourself.
Where the span comes from
k is the number of rows in each period. When you omit it, k is input.n_slots. stage_sizes must multiply to k.
nb_cols is the number of sums. The default is 1, one sum of the period. stage_sizes must multiply to n_slots / nb_cols.
HomExpand takes no k. It expands the input.n_slots values already set on the input. Set that with HomParams.n_slots for the primary input, or with pipeline custom_n_slots for a named input.
HomAdd, HomMul, and the reshapes never work across slots. They take no stage plan.
# HomParams.n_slots = 128 HomSumSlots() # nb_cols defaults to 1: one sum of the period HomRunningSum() # k defaults to n_slots, not the whole ring HomExpand(expand_axis=1, ...) # width is input.n_slots HomMul(...) # element-wise: no span, no stage plan
n_slots, and HomExpand reads it from the ciphertext. The compiler checks the match at compile time, but picking the right number is on you.The price of a stage plan
Wide stages need many distinct offsets, and every offset is one more entry in the rotation key: a bigger key to generate and upload, more GPU memory on the worker. In exchange, fewer rounds: the query runs faster.
Narrow stages reuse a handful of offsets, so the key stays small and cheap to upload. In exchange, more rounds: each round is one more pass over the ciphertext, so the query is somewhat slower.
Stages spend levels
HomExpand and HomRunningSum, run a mod-switch as they go. How often is set by stages_per_level: how many stages to group together before each mod-switch.1 (default)Mod-switch after every stagelen(stage_sizes)One mod-switch, at the endlen(stage_sizes) / stages_per_level. A default plan for a span of 1024, left at stages_per_level = 1, spends ten levels. More stages spend more levels unless you group the mod-switches.HomRunningSum defaults k to n_slots, the period, not the whole ring. A smaller period uses a shorter plan and fewer levels.HomSumSlots is the exception. It never mod-switches between stages, has no stages_per_level parameter, and spends zero levels whatever its stage plan. Its only costs are key size and rounds.
Operators that take a stage plan
HomExpandOne output per value in the period. The slot axis must be exactly n_slots long.HomRunningSumRunning sum inside each period. k is the row count and defaults to n_slots.HomSumSlotsSums nb_cols columns of the period. Zero padding is allowed. Spends no levels.HomMatMul runs a slot sum when the axis it multiplies over is the n axis. Encryption already repeats the n_slots block for you, so there is no client-side Repeat operator.Same name, different thing: Bootstrap has its own internal transforms. Those are separate from the stage plans on this page; nothing here applies to them.
Choosing a stage plan
While developing, leave stage_sizes unset. The log-depth plan keeps the evaluation key at its smallest, and key generation and upload stay quick.
If key size, upload time, and worker GPU memory are acceptable, widen the stages to cut rounds and speed up the query.
Keep the stages, raise stages_per_level. You keep the small key and pay fewer levels; only the round count stays.
# Default: log-depth tree (smallest key) HomExpand() # input.n_slots=128: 7 stages, 7 levels # Collapsed: one wide stage (bigger key, fastest) HomExpand(stage_sizes=(128,)) # 1 stage, 1 level # Small key, grouped mod-switches (fewest levels for this tree) HomExpand(stage_sizes=(2,)*7, stages_per_level=7) # 7 stages, 1 level