torch_to_nnef.op.extras

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.

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.