yobx.torch

main functions

• yobx.torch.to_onnx

modules

• yobx.torch._model_eval_cases
• yobx.torch.export_helper
• yobx.torch.export_options
• yobx.torch.fake_tensor_helper
• yobx.torch.flatten
• yobx.torch.flatten_helper
• yobx.torch.input_observer
• yobx.torch.model_eval_cases
• yobx.torch.patch
• yobx.torch.tiny_models
• yobx.torch.torch_helper
• yobx.torch.tracing
• yobx.torch.validate
• yobx.torch.in_torch
• yobx.torch.in_transformers
• yobx.torch.interpreter

register_flattening_functions

yobx.torch.register_flattening_functions(patch_transformers: bool = False, verbose: int = 0) → Generator[Callable, None, None]

The context manager registers flattening functions the exporter needs to handle any custom class. This is used to create a signature taking only tensors as inputs even though the code shows nested structures.

from yobx.torch import register_flattening_functions

with register_flattening_functions(patch_transformers=True):
    # ...