hetorch.passes.builtin.relinearization
RelinearizationInsertionPass: Insert relinearization operations after multiplications (CKKS)
Classes
RelinearizationInsertionPass(strategy: <class 'str'>)
Insert relinearization operations after ciphertext multiplications for CKKS scheme
In CKKS (and other schemes), ciphertext-ciphertext multiplications increase the ciphertext size (number of polynomials). Relinearization reduces the ciphertext back to its original size using relinearization keys.
This pass supports two strategies:
- "eager": Relinearize after every ciphertext multiplication
- "lazy": Only relinearize when necessary (before next multiplication or at output)
Attributes: strategy: Relinearization strategy ("eager" or "lazy")
Methods:
__init__(self, strategy: str = 'lazy')
Initialize RelinearizationInsertionPass
Args: strategy: "eager" (relin after each cmult) or "lazy" (relin only when needed)
analyze_cost(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> hetorch.backend.cost_model.CostAnalysis
Analyze cost impact of this pass
Args: graph_module: Graph module to analyze context: Compilation context
Returns: Cost analysis result
transform(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> torch.fx.graph_module.GraphModule
Apply relinearization insertion transformation
Args: graph_module: Input graph module context: Compilation context
Returns: Transformed graph module with relinearization operations
validate(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> bool
Validate that the pass can be applied
Args: graph_module: Graph module to validate context: Compilation context
Returns: True if validation passes
Raises: PassValidationError: If validation fails