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.
2.7 KiB
Supporting New ONNX Nodes in Torus ML
Torus ML supports a wide range of models through the integration of ONNX nodes. In case a specific ONNX node is missing, developers need to add support for the new ONNX nodes.
Operator Implementation
Floating-point Implementation
The ops_impl.py file is responsible for implementing the computation of ONNX operators using floating-point arithmetic. The implementation should mirror the behavior of the corresponding ONNX operator precisely. This includes adhering to the expected inputs, outputs, and operational semantics.
Refer to the ONNX documentation to grasp the expected behavior, inputs and outputs of the operator.
Operator Mapping
After implementing the operator in ops_impl.py, you need to import it into onnx_utils.py and map it within the ONNX_OPS_TO_NUMPY_IMPL dictionary. This mapping is crucial for the framework to recognize and utilize the new operator.
Quantized Operator
Quantized operators are defined in quantized_ops.py and are used to handle integer arithmetic. Their implementation is required for the new ONNX to be executed in FHE.
There exist two types of quantized operators:
- Univariate Non-Linear Operators: Such operator applies transformation on every element of the input without changing its shape. Sigmoid, Tanh, ReLU are examples of such operation. The sigmoid in this file is simply supported as follows:
class QuantizedSigmoid(QuantizedOp):
"""Quantized sigmoid op."""
_impl_for_op_named: str = "Sigmoid"
- Linear Layers: Linear layers like
GemmandConvrequire specific implementations for integer arithmetic. Please refer to theQuantizedGemmandQuantizedConvimplementations for reference.
Adding Tests
Proper testing is essential to ensure the correctness of the new ONNX node support.
There are many locations where tests can be added:
test_onnx_ops_impl.py: Tests the implementation of the ONNX node in floating points.test_quantized_ops.py: Tests the implementation of the ONNX node in integer arithmetic.- Optional:
test_compile_torch.py: Tests the implementation of a specific torch model that contains the new ONNX operator. The model needs to be added intorch_models.py.
Update Documentation
Finally, update the documentation to reflect the newly supported ONNX node.