torch_to_nnef.torch_graph.ir_op

Abstractions used in torch_to_nnef internal graph Operations IR.

(decoupled from PyTorch and NNEF)

The goal is that these elements are: - extracted/parsed from PyTorch graph operations - translated to NNEF graph operations

Classes:

Name	Description
InputsAlignBetweenAtenAndTorch	Mapping inputs between Python torch.$1 and cpp aten::$2.
TorchOp

InputsAlignBetweenAtenAndTorch

Mapping inputs between Python torch.$1 and cpp aten::$2.

Because function arguments are not 1 to 1

Methods:

Name	Description
aten_logspace	aten::logspace(start, end, steps, base, *, dtype, ...).
aten_trilu_indices	aten::{tril,triu}_indices(row, col, offset, *, dtype, ...).
aten_vander	aten::vander(self, N, increasing).
aten_window_constant	aten::{bartlett,hann,hamming,blackman,kaiser}_window.

aten_logspace staticmethod

aten_logspace(args, kwargs)

aten::logspace(start, end, steps, base, *, dtype, ...).

First four are positional; the rest are keyword-only on the aten op.

aten_trilu_indices staticmethod

aten_trilu_indices(args, kwargs)

aten::{tril,triu}_indices(row, col, offset, *, dtype, ...).

Only row / col / offset are positional; the rest are keyword-only on the aten op. Trim the tail and lift dtype (kept; layout / device / pin_memory are dropped since trace- time evaluation doesn't need them).

aten_vander staticmethod

aten_vander(args, kwargs)

aten::vander(self, N, increasing).

N may be None for the default (= len(self)); torch's signature accepts None directly so we forward as-is.

aten_window_constant staticmethod

aten_window_constant(args, kwargs)

aten::{bartlett,hann,hamming,blackman,kaiser}_window.

Positional: window_length (+ periodic for the non-Bartlett variants); rest is keyword-only on the aten op.

TorchOp dataclass

TorchOp(kind: str, module_path: str, inputs: List[Data], outputs: List[TtupleOrVar], scope: str, op_ref: Optional[Callable], call_name: Optional[str])

Methods:

Name	Description
call_op	Produce operation output based on traced inputs with real torch call.
output_aliases_constant_input	Detect view-like results derived from constant inputs.
realise_output_type_and_size	Trace output and try to find type shape and constant realisation.
update_call_op_arg_kwargs	Custom adaptation to call aten fn with torch exposed py fn.

call_op

call_op()

Produce operation output based on traced inputs with real torch call.

This operation call is done via self.args arguments (for now). Which means that we need to have all args needed in parameters order, following at least 1 underling torch operation signature.

NOTE: we use a different approach than original torch.onnx which pass parameter by keyword arguments, this is due to the fact that we are not aware of argument name being provided in exported graph ( from what we understand torch.onnx solve this via explicit rerouting of all signatures, which might be a bit bulky in most case ).

output_aliases_constant_input

output_aliases_constant_input(result: Tensor) -> bool

Detect view-like results derived from constant inputs.

These outputs should keep shape/dtype metadata, but should not be recorded as new constants. Otherwise a graph like Granite MoE's packed expert weight selects serializes every expert slice as an independent tensor in addition to the original packed parameter.

realise_output_type_and_size

realise_output_type_and_size(approx: bool = True) -> bool

Trace output and try to find type shape and constant realisation.

update_call_op_arg_kwargs

update_call_op_arg_kwargs(args)

Custom adaptation to call aten fn with torch exposed py fn.