zeekay
e037e8f7ea
The distribution called itself torus-ml while shipping a package directory named concrete, so the name a user installs and the name they import disagreed — and most of the tree already said torus.ml, which could not resolve against src/concrete. Moved src/concrete to src/torus and pointed the include at it. The dependency was Zama's concrete-python from their index. It is torus-fhe from ours, which is what luxfhe/torus already ships: that frontend declares torus-fhe and ships a torus/fhe module, so every 'from torus.fhe' in this tree was already naming our compiler and only the dependency lagged. Two distributions share one namespace, which is the shape the upstream this forked from used and the reason the split works: torus-fhe gives torus.fhe, torus-ml gives torus.ml. Swept every other spelling — identifiers, enum members, workflow variables, filenames, notebooks. LICENSE and NOTICE keep the upstream attribution a fork owes; the figures carry the old name in JPEG metadata a text sweep cannot reach.
52 KiB
module torus.ml.pytest.torch_models
Torch modules for our pytests.
class MultiOutputModel
Multi-output model.
method __init__
__init__() → None
Torch Model.
method forward
forward(x, y)
Forward pass.
Args:
x(torch.Tensor): The input of the model.y(torch.Tensor): The input of the model.
Returns:
Tuple[torch.Tensor. torch.Tensor]: Outputs of the network.
class SimpleNet
Fake torch model used to generate some onnx.
method __init__
__init__() → None
method forward
forward(inputs)
Forward function.
Arguments:
inputs: the inputs of the model.
Returns:
torch.Tensor: the result of the computation
class FCSmall
Torch model for the tests.
method __init__
__init__(input_output, activation_function, hidden=None)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class FC
Torch model for the tests.
method __init__
__init__(activation_function, input_output=3072)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class CNN
Torch CNN model for the tests.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class CNNMaxPool
Torch CNN model for the tests with a max pool.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class CNNOther
Torch CNN model for the tests.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class CNNInvalid
Torch CNN model for the tests.
method __init__
__init__(activation_function, groups)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class CNNGrouped
Torch CNN model with grouped convolution for compile torch tests.
method __init__
__init__(input_output, activation_function, groups)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class NetWithLoops
Torch model, where we reuse some elements in a loop.
Torch model, where we reuse some elements in a loop in the forward and don't expect the user to define these elements in a particular order.
method __init__
__init__(activation_function, input_output, n_fc_layers)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class MultiInputNN
Torch model to test multiple inputs forward.
method __init__
__init__(input_output, activation_function)
method forward
forward(x, y)
Forward pass.
Args:
x: the first input of the NNy: the second input of the NN
Returns: the output of the NN
class MultiInputNNConfigurable
Torch model to test multiple inputs forward.
method __init__
__init__(use_conv, use_qat, input_output, n_bits)
method forward
forward(x, y)
Forward pass.
Args:
x: the first input of the NNy: the second input of the NN
Returns: the output of the NN
class MultiInputNNDifferentSize
Torch model to test multiple inputs with different shape in the forward pass.
method __init__
__init__(
input_output,
activation_function=None,
is_brevitas_qat=False,
n_bits=3
)
method forward
forward(x, y)
Forward pass.
Args:
x: The first input of the NN.y: The second input of the NN.
Returns: The output of the NN.
class BranchingModule
Torch model with some branching and skip connections.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class BranchingGemmModule
Torch model with some branching and skip connections.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class UnivariateModule
Torch model that calls univariate and shape functions of torch.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class StepActivationModule
Torch model implements a step function that needs Greater, Cast and Where.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass with a quantizer built into the computation graph.
Args:
x: the input of the NN
Returns: the output of the NN
class NetWithConcatUnsqueeze
Torch model to test the concat and unsqueeze operators.
method __init__
__init__(activation_function, input_output, n_fc_layers)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class MultiOpOnSingleInputConvNN
Network that applies two quantized operations on a single input.
method __init__
__init__(can_remove_input_tlu: bool)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class FCSeq
Torch model that should generate MatMul->Add ONNX patterns.
This network generates additions with a constant scalar
method __init__
__init__(input_output, act)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class FCSeqAddBiasVec
Torch model that should generate MatMul->Add ONNX patterns.
This network tests the addition with a constant vector
method __init__
__init__(input_output, act)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN
class TinyCNN
A very small CNN.
method __init__
__init__(n_classes, act) → None
Create the tiny CNN with two conv layers.
Args:
n_classes: number of classesact: the activation
method forward
forward(x)
Forward the two layers with the chosen activation function.
Args:
x: the input of the NN
Returns: the output of the NN
class TinyQATCNN
A very small QAT CNN to classify the sklearn digits data-set.
This class also allows pruning to a maximum of 10 active neurons, which should help keep the accumulator bit-width low.
method __init__
__init__(
n_classes,
n_bits,
n_active,
signed,
narrow,
power_of_two_scaling
) → None
Construct the CNN with a configurable number of classes.
Args:
n_classes(int): number of outputs of the neural netn_bits(int): number of weight and activation bits for quantizationn_active(int): number of active (non-zero weight) neurons to keepsigned(bool): whether quantized integer values are signednarrow(bool): whether the range of quantized integer values is narrow/symmetricpower_of_two_scaling(bool): whether to use power-of-two scaling quantizers which allows to test the round PBS optimization when the scales are power-of-two
method forward
forward(x)
Run inference on the tiny CNN, apply the decision layer on the reshaped conv output.
Args:
x: the input to the NN
Returns: the output of the NN
method toggle_pruning
toggle_pruning(enable)
Enable or remove pruning.
Args:
enable: if we enable the pruning or not
class StepFunctionPTQ
Torch model implements a step function that needs Greater, Cast and Where.
method __init__
__init__(input_output, activation_function, n_bits=2, disable_bit_check=False)
method forward
forward(x)
Forward pass with a quantizer built into the computation graph.
Args:
x: the input of the NN
Returns: the output of the NN
class QATTestModule
Torch model that implements a simple non-uniform quantizer.
method __init__
__init__(activation_function)
method forward
forward(x)
Forward pass with a quantizer built into the computation graph.
Args:
x: the input of the NN
Returns: the output of the NN
class SingleMixNet
Torch model that with a single conv layer that produces the output, e.g., a blur filter.
method __init__
__init__(use_conv, use_qat, inp_size, n_bits)
method forward
forward(x)
Execute the single convolution.
Args:
x: the input of the NN
Returns: the output of the NN
class DoubleQuantQATMixNet
Torch model that with two different quantizers on the input.
Used to test that it keeps the input TLU.
method __init__
__init__(use_conv, use_qat, inp_size, n_bits)
method forward
forward(x)
Execute the single convolution.
Args:
x: the input of the NN
Returns: the output of the NN
class TorchSum
Torch model to test the ReduceSum ONNX operator in a leveled circuit.
method __init__
__init__(dim=(0,), keepdim=True)
Initialize the module.
Args:
dim(Tuple[int]): The axis along which the sum should be executedkeepdim(bool): If the output should keep the same dimension as the input or not
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model
Returns:
torch_sum(torch.tensor): The sum of the input's tensor elements along the given axis
class NetWithConstantsFoldedBeforeOps
Torch QAT model that does not quantize the inputs.
method __init__
__init__(
hparams: dict,
bits: int,
act_quant=<class 'brevitas.quant.scaled_int.Int8ActPerTensorFloat'>,
weight_quant=<class 'brevitas.quant.scaled_int.Int8WeightPerTensorFloat'>
)
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model
Returns:
torch.tensor: Output of the network
class ShapeOperationsNet
Torch QAT model that reshapes the input.
method __init__
__init__(is_qat)
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model
Returns:
torch.tensor: Output of the network
class PaddingNet
Torch QAT model that applies various padding patterns.
method __init__
__init__()
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model
Returns:
torch.tensor: Output of the network
class QuantCustomModel
A small quantized network with Brevitas, trained on make_classification.
method __init__
__init__(
input_shape: int,
output_shape: int,
hidden_shape: int = 100,
n_bits: int = 5,
act_quant=<class 'brevitas.quant.scaled_int.Int8ActPerTensorFloat'>,
weight_quant=<class 'brevitas.quant.scaled_int.Int8WeightPerTensorFloat'>,
bias_quant=None
)
Quantized Torch Model with Brevitas.
Args:
input_shape(int): Input sizeoutput_shape(int): Output sizehidden_shape(int): Hidden sizen_bits(int): Bit of quantizationweight_quant(brevitas.quant): Quantization protocol of weightsact_quant(brevitas.quant): Quantization protocol of activations.bias_quant(brevitas.quant): Quantizer for the linear layer bias
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model.
Returns:
torch.tensor: Output of the network.
class TorchCustomModel
A small network with Brevitas, trained on make_classification.
method __init__
__init__(input_shape, hidden_shape, output_shape)
Torch Model.
Args:
input_shape(int): Input sizeoutput_shape(int): Output sizehidden_shape(int): Hidden size
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model.
Returns:
torch.tensor: Output of the network.
class ConcatFancyIndexing
Concat with fancy indexing.
method __init__
__init__(
input_shape,
hidden_shape,
output_shape,
n_bits: int = 4,
n_blocks: int = 3
) → None
Torch Model.
Args:
input_shape(int): Input sizeoutput_shape(int): Output sizehidden_shape(int): Hidden sizen_bits(int): Number of bitsn_blocks(int): Number of blocks
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model.
Returns:
torch.tensor: Output of the network.
class PartialQATModel
A model with a QAT Module.
method __init__
__init__(input_shape: int, output_shape: int, n_bits: int)
method forward
forward(x)
Forward pass.
Args:
x(torch.tensor): The input of the model.
Returns:
torch.tensor: Output of the network.
class EncryptedMatrixMultiplicationModel
PyTorch module for performing matrix multiplication between two encrypted values.
method __init__
__init__(input_output, activation_function)
method forward
forward(input1)
Forward function for matrix multiplication.
Args:
input1(torch.Tensor): The first input tensor.
Returns:
torch.Tensor: The result of the matrix multiplication.
class ManualLogisticRegressionTraining
PyTorch module for performing SGD training.
method __init__
__init__(learning_rate=0.1)
method forward
forward(x, y, weights, bias)
Forward function for matrix multiplication.
Args:
x(torch.Tensor): The training data tensor.y(torch.Tensor): The target tensor.weights(torch.Tensor): The weights to be learned.bias(torch.Tensor): The bias to be learned.
Returns:
torch.Tensor: The updated weights after performing a training step.
method predict
predict(x, weights, bias)
Predicts based on weights and bias as inputs.
Args:
x(torch.Tensor): Input data tensor.weights(torch.Tensor): Weights tensor.bias(torch.Tensor): Bias tensor.
Returns:
torch.Tensor: The predicted outputs for the given inputs.
class WhereNet
Simple network with a where operation for testing.
method __init__
__init__(n_hidden)
method forward
forward(x)
Forward pass.
Args:
x(torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying the where operation.
class AddNet
Torch model that performs a simple addition between two inputs.
method __init__
__init__(use_conv, use_qat, input_output, n_bits)
method forward
forward(x, y)
Forward pass.
Args:
x: First input tensor.y: Second input tensor.
Returns: Result of adding x and y.
class ExpandModel
Minimalist network that expands the input tensor to a larger size.
method __init__
__init__(is_qat)
method forward
forward(x)
Expand the input tensor to the target size.
Args:
x(torch.Tensor): Input tensor.
Returns:
torch.Tensor: Expanded tensor.
class Conv1dModel
Small model that uses a 1D convolution operator.
method __init__
__init__(input_output, activation_function) → None
method forward
forward(x)
Forward pass.
Args:
x(torch.Tensor): The model's input.
Returns:
torch.Tensor: The model's output.
class IdentityExpandModel
Model that only adds an empty dimension at axis 0.
This model is mostly useful for testing the composition feature.
method forward
forward(x)
Forward pass.
Args:
x(torch.Tensor): The input of the model.
Returns:
Tuple[torch.Tensor. torch.Tensor]: Outputs of the network.
class IdentityExpandMultiOutputModel
Model that only adds an empty dimension at axis 0, and returns the initial input as well.
This model is mostly useful for testing the composition feature.
method forward
forward(x)
Forward pass.
Args:
x(torch.Tensor): The input of the model.
Returns:
Tuple[torch.Tensor. torch.Tensor]: Outputs of the network.
class TorchDivide
Torch model that performs a encrypted division between two inputs.
method __init__
__init__(input_output, activation_function)
method forward
forward(x, y)
Forward pass.
Args:
x(torch.Tensor): The first input tensor.y(torch.Tensor): The second input tensor.
Returns:
torch.Tensor: The result of the division.
class TorchMultiply
Torch model that performs a encrypted multiplication between two inputs.
method __init__
__init__(input_output, activation_function)
method forward
forward(x, y)
Forward pass.
Args:
x(torch.Tensor): The first input tensor.y(torch.Tensor): The second input tensor.
Returns:
torch.Tensor: The result of the multiplication.
class EmbeddingModel
A torch model with an embedding layer.
method __init__
__init__(
num_embeddings,
embedding_dim,
activation_function=<class 'torch.nn.modules.activation.ReLU'>
)
method forward
forward(x)
Forward pass.
Args:
x(torch.Tensor): The input tensor containing indices.
Returns:
torch.Tensor: The output tensor after embedding.
class AllZeroCNN
A CNN class that has all zero weights and biases.
method __init__
__init__(input_output, activation_function)
method forward
forward(x)
Forward pass.
Args:
x: the input of the NN
Returns: the output of the NN