ir_module_tracer

torch_to_nnef.torch_graph.ir_module_tracer

implements :class:TorchModuleTracer.

It traces a torch.nn.Module with torch.jit.trace and exposes the resulting graph.

with few related helper functions (e.g. _is_io_quantized_module and maybe_quantize_args_tensor) provide small utilities used during tracing.

TorchModuleTracer

TorchModuleTracer(module: nn.Module, traced_module: T.Optional[torch.jit.TracedModule] = None, fn_name: str = 'forward', args: T.Optional[T.Tuple[T.Any, ...]] = None)

Evaluate Optimized traced Function code so that signature always match.

original Module is passed to do proper un-boxing later on. This is needed because we have a re-routing based on actual module classtype.

Create a tracer for module.

The tracer stores the original module, an optional pre‑traced torch.jit.TracedModule (which allows re‑use of a previously computed trace), the name of the forward method to trace, and the arguments used for tracing. The arguments are post‑processed by :func:maybe_quantize_args_tensor to ensure compatibility with quantized modules.

torch_graph property

torch_graph

Return the underlying PyTorch graph object.

The actual torch.Graph is retrieved from the traced module. When a different forward method is requested (fn_name differs from "forward"), the corresponding sub‑graph is returned instead.

traced_module property

traced_module

Return the traced module, computing it lazily if required.

If self._traced_module is None the method will perform a jit.trace on self.mod with self.args while handling possible PyTorch version nuances. Any RuntimeError raised by torch.jit.trace is wrapped into a :class:~torch_to_nnef.exceptions.T2NErrorTorchJitTraceFailed exception.

maybe_quantize_args_tensor

maybe_quantize_args_tensor(module, args)

Quantize tensors in args if module expects quantized input.

The function walks the args tuple, and for each tensor that is not already quantized, it creates a fake quantized representation using torch.quantize_per_tensor with a dummy scale/zero point. The quantized tensor is not representative of real data; it is solely used to keep the tracing machinery happy when the module expects quantized inputs.

Parameters:

Name	Type	Description	Default
module		A torch.nn.Module instance.	required
args		A tuple of inputs supplied to the module. None is accepted for modules that do not require any positional arguments.	required

Returns:

Type	Description
	A potentially modified tuple where tensors are quantized when
	necessary.