Skip to main content

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