torch_to_nnef.op.helper

Classes:

Name	Description
`OpHelper`
`OpRegistry`

Functions:

Name	Description
`add_tensor_variable_node_as_nnef_tensor`	Create NNEF tensor and register in graph from torch_graph.Data node.
`cast_and_add_nnef_operation`	Ensure to cast parameters before adding operation to NNEF graph.
`cast_inputs_and_attrs`	Catch input or attr that would still be torch_graph values into NNEF.
`cast_to_if_not_dtype_and_variable`	Force casting not expressed in IR graph in case of div for example.
`maybe_align_inputs_ranks`	Ensure consistent rank between inputs and outputs with regard to spec.
`pick_axis`	Enforce that axis, axes ect does contains only positive values.
`pick_index_in_axis`	Enforce that index in axis does contains only values within bounds.
`resolve_attr_axis_size`	Resolve `input_node.shape[axis]` for use in a NNEF op attribute.

OpHelper

OpHelper(g, node, name_to_tensor, null_ref, inference_target)

Methods:

Name	Description
`add_intermediate_op`	Emit an op whose output is a fresh intermediate NTensor.
`add_single_output_op_from_ir_datas`	Use input_nodes Data instead of nnef.Tensor.
`emit_dtype_cast`	Emit a shape-preserving `tract_core_cast` for one operand.

add_intermediate_op

add_intermediate_op(src: Tensor, op_type: str, attrs: Optional[Dict[str, Any]], new_shape: Sequence[int], suffix: str) -> NTensor

Emit an op whose output is a fresh intermediate NTensor.

add_single_output_op derives its output NNEF tensor from node.outputs[0], which for chained-helper emits means every intermediate would inherit the final op's shape. For multi- step decomposition we need to author the intermediate NTensor by hand: this helper does that.

The intermediate inherits src.dtype (no dtype-change ops like cast go through this helper -- use the dedicated cast emitter for those).

Parameters:

Name	Type	Description	Default
`src`	`Tensor`	the upstream NNEF tensor feeding this op.	required
`op_type`	`str`	NNEF op type to emit.	required
`attrs`	`Optional[Dict[str, Any]]`	op attributes (or None).	required
`new_shape`	`Sequence[int]`	shape of the intermediate output.	required
`suffix`	`str`	appended to `src.name` to name the new tensor.	required

Returns:

Type	Description
`Tensor`	The newly-created NNEF tensor.

add_single_output_op_from_ir_datas

add_single_output_op_from_ir_datas(nnef_op_type: str, input_nodes: List[Data], output_tensor_name_suffix: str = '', force_full_output_tensor_name: str = '', reuse_if_name_exists: bool = False, **kwargs) -> TorchOp

Use input_nodes Data instead of nnef.Tensor.

Also nnefe

emit_dtype_cast

emit_dtype_cast(node, tensor, to_tract_dtype, full_name=None)

Emit a shape-preserving tract_core_cast for one operand.

A dtype cast never changes shape, but add_single_output_op shapes the emitted op's output from node.outputs[0] (the parent op's result). When the cast is an auxiliary op (e.g. an implicit dtype-promotion inserted for a binary op, or a bool/branch cast for a select) whose operand has a lower rank than the parent output, that borrowed shape fools the rank aligner (:func:maybe_align_inputs_ranks) into thinking the operand is already full-rank, so it skips the broadcast unsqueeze and the operand mis-broadcasts under tract's left-aligned expansion (torch right-aligns). Restore the operand's own shape so every emitted cast is faithful regardless of the parent node.

OpRegistry

OpRegistry(torch_mod_id: str)

Methods:

Name	Description
`register`	By default we take the name of the function if not specified.

register

register(torch_op_ids: Optional[List[str]] = None)

By default we take the name of the function if not specified.

add_tensor_variable_node_as_nnef_tensor

add_tensor_variable_node_as_nnef_tensor(g: Graph, node: TensorVariable, name_to_tensor: Dict[str, Tensor], name_suffix: str = '', prevent_variable: bool = False, force_full_output_tensor_name: Optional[str] = None) -> NTensor

Create NNEF tensor and register in graph from torch_graph.Data node.

It automatically adds variable if node is a torch tensor is associated (it avoids bloating nnef graph file with matrix values)

cast_and_add_nnef_operation

cast_and_add_nnef_operation(name_to_tensor: str, **kwargs)

Ensure to cast parameters before adding operation to NNEF graph.

cast_inputs_and_attrs

cast_inputs_and_attrs(inputs, attrs, g, name_to_tensor)

Catch input or attr that would still be torch_graph values into NNEF.

cast_to_if_not_dtype_and_variable

cast_to_if_not_dtype_and_variable(g, name_to_tensor, node, nnef_tensor: Tensor, cast_to: dtype, suffix: str = '')

Force casting not expressed in IR graph in case of div for example.

This is neccessary since tract and maybe other inference engine may not cast implicitly to float during div operation for example leading to rounding issues.

maybe_align_inputs_ranks

maybe_align_inputs_ranks(g: Graph, inputs: Sequence[Tensor], outputs: Sequence[Tensor], op_type: str) -> T.Sequence[NTensor]

Ensure consistent rank between inputs and outputs with regard to spec.

May unsqueeze at 0 rank n time to align inputs

This is done at export time and not inference time because: - inference implementation may use 1 dim expansion from left to right like Tract or Tensorflow instead of PyTorch expansion which happen in opposite direction.

pick_axis

pick_axis(input_node, rank: int) -> int

Enforce that axis, axes ect does contains only positive values.

pick_index_in_axis

pick_index_in_axis(input_node, rank: int, index: int, check_is_positive: bool = True) -> int

Enforce that index in axis does contains only values within bounds.

Because in case of tract out of bound is not supported !

resolve_attr_axis_size

resolve_attr_axis_size(op_helper, input_node, axis: int)

Resolve input_node.shape[axis] for use in a NNEF op attribute.

When inference_target.has_dynamic_axes, emit (and cache) a tract_core_shape_of -> slice -> squeeze chain to extract the runtime size of axis and return an nnef.Identifier referencing that scalar. Otherwise, return the static int(input_node.shape[axis]).

Use this anywhere a shape-derived value lands in an op attribute (e.g. reshape(shape=[...]), tile(repeats=[...]), slice(end=...)) so the same emitter works in both static and dyn-axes modes without baking the trace-time size into the exported graph.