5.1 KiB
module torus.ml.torch.hybrid_backprop_linear
Linear layer implementations for backprop FHE-compatible models.
class ForwardModuleLinear
Forward module for linear layers.
method __init__
__init__(weight, bias=None, weight_transposed=False)
method forward
forward(input_tensor)
Forward pass for linear layers.
Args:
input_tensor: The input tensor.
Returns: The output tensor after applying the linear transformation.
class BackwardModuleLinear
Backward module for linear layers.
method __init__
__init__(weight, weight_transposed=False)
method forward
forward(grad_output)
Backward pass for linear layers.
Args:
grad_output: The gradient output tensor.
Returns: The gradient input tensor after applying the backward pass.
class CustomLinear
Custom linear module.
method __init__
__init__(weight, bias=None, weight_transposed=False)
method forward
forward(input_tensor)
Forward pass of the custom linear module.
Args:
input_tensor: The input tensor.
Returns: The output tensor after applying the custom linear module.
class ForwardBackwardModule
Custom autograd function for forward and backward passes.
method backward
backward(ctx, grad_output)
Backward pass of the custom autograd function.
Args:
ctx: The context object.grad_output: The gradient output tensor.
Returns: The gradient input tensor after applying the backward pass.
method forward
forward(ctx, input_tensor, forward_module, backward_module)
Forward pass of the custom autograd function.
Args:
ctx: The context object.input_tensor: The input tensor.forward_module: The forward module.backward_module: The backward module.
Returns: The output tensor after applying the forward pass.