torch_to_nnef.op.extras

Dispatch for custom PyTorch ops registered under the t2n_extra:: namespace.

Users can ship their own opaque ops by declaring them via torch.library.custom_op("t2n_extra::<name>", ...) and providing a matching handler here. The handler emits a NNEF op (typically a custom fragment), so the trace stays a single opaque node end-to-end instead of being decomposed by torch.jit.trace.

Today this powers Mamba's selective scan: see torch_to_nnef.op.extras.scan_ops:ssm_scan for the wiring pattern.

Modules:

Name	Description
exp_norm	Handlers for t2n_extra::exp_unit_norm / exp_mean_norm ops.
scan_ops	Handlers for t2n_extra::* scan-shaped ops.

Functions:

Name	Description
register	Decorator: register a handler for t2n_extra::<op_name>.
t2n_extra_to_nnef_tensor_and_ops	Dispatch a t2n_extra::* op node to its registered handler.

register

register(op_name: str)

Decorator: register a handler for t2n_extra::<op_name>.

t2n_extra_to_nnef_tensor_and_ops

t2n_extra_to_nnef_tensor_and_ops(g, node, name_to_tensor, null_ref, *, torch_graph, inference_target) -> T.List[str]

Dispatch a t2n_extra::* op node to its registered handler.