ml-sdk/docs/references/api/concrete.ml.pytest.torch_models.md

module concrete.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 NN
• y: 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 NN
• y: 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 classes
• act: 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 net
• n_bits (int): number of weight and activation bits for quantization
• n_active (int): number of active (non-zero weight) neurons to keep
• signed (bool): whether quantized integer values are signed
• narrow (bool): whether the range of quantized integer values is narrow/symmetric
• power_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 executed
• keepdim (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 size
• output_shape (int): Output size
• hidden_shape (int): Hidden size
• n_bits (int): Bit of quantization
• weight_quant (brevitas.quant): Quantization protocol of weights
• act_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 size
• output_shape (int): Output size
• hidden_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 size
• output_shape (int): Output size
• hidden_shape (int): Hidden size
• n_bits (int): Number of bits
• n_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