Files
torus-ml/docs/references/api/torus.ml.onnx.onnx_model_manipulations.md
zeekay e037e8f7ea torus: one name for the package, and it is ours
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.
2026-08-27 13:20:50 -07:00

5.8 KiB

module torus.ml.onnx.onnx_model_manipulations

Some code to manipulate models.


function simplify_onnx_model

simplify_onnx_model(onnx_model: ModelProto)

Simplify an ONNX model, removes unused Constant nodes and Identity nodes.

Args:

  • onnx_model (onnx.ModelProto): the model to simplify.

function remove_unused_constant_nodes

remove_unused_constant_nodes(onnx_model: ModelProto)

Remove unused Constant nodes in the provided onnx model.

Args:

  • onnx_model (onnx.ModelProto): the model for which we want to remove unused Constant nodes.

function remove_identity_nodes

remove_identity_nodes(onnx_model: ModelProto)

Remove identity nodes from a model.

Args:

  • onnx_model (onnx.ModelProto): the model for which we want to remove Identity nodes.

function keep_following_outputs_discard_others

keep_following_outputs_discard_others(
    onnx_model: ModelProto,
    outputs_to_keep: Iterable[str]
)

Keep the outputs given in outputs_to_keep and remove the others from the model.

Args:

  • onnx_model (onnx.ModelProto): the ONNX model to modify.
  • outputs_to_keep (Iterable[str]): the outputs to keep by name.

function remove_node_types

remove_node_types(onnx_model: ModelProto, op_types_to_remove: List[str])

Remove unnecessary nodes from the ONNX graph.

Args:

  • onnx_model (onnx.ModelProto): The ONNX model to modify.
  • op_types_to_remove (List[str]): The node types to remove from the graph.

Raises:

  • ValueError: Wrong replacement by an Identity node.

function clean_graph_at_node_op_type

clean_graph_at_node_op_type(
    onnx_model: ModelProto,
    node_op_type: str,
    fail_if_not_found: bool = True
)

Remove the first node matching node_op_type and its following nodes from the ONNX graph.

Args:

  • onnx_model (onnx.ModelProto): The onnx model.
  • node_op_type (str): The node's op_type whose following nodes as well as itself will be removed.
  • fail_if_not_found (bool): If true, raise an error if no node matched the given op_type.

Raises:

  • ValueError: If no node matched the given op_type and fail_if_not_found is set to True

function clean_graph_after_node_op_type

clean_graph_after_node_op_type(
    onnx_model: ModelProto,
    node_op_type: str,
    fail_if_not_found: bool = True
)

Remove the nodes following first node matching node_op_type from the ONNX graph.

Args:

  • onnx_model (onnx.ModelProto): The onnx model.
  • node_op_type (str): The node's op_type whose following nodes will be removed.
  • fail_if_not_found (bool): If true, raise an error if no node matched the given op_type.

Raises:

  • ValueError: If no node matched the given op_type and fail_if_not_found is set to True

function convert_first_gather_to_matmul

convert_first_gather_to_matmul(
    onnx_model: ModelProto
) → Tuple[Union[ModelProto, NoneType], ModelProto]

Convert the first Gather node to a matrix multiplication node.

In FHE, Gather is a costly operation since it can involve many PBS. When it appears first in the onnx model, we can remove it and replace it by a matrix multiplication node by converting the indices to a one-hot encoding.

Args:

  • onnx_model (onnx.ModelProto): The onnx model.

Returns:

  • Tuple[Optional[onnx.ModelProto], onnx.ModelProto]: The pre-processing model and the modified onnx model.