yobx.torch.patch#

yobx.torch.patch.apply_patches_for_model(patch_torch: bool = False, patch_transformers: bool = False, verbose: int = 0, model: Module | None = None, extra_patches: Iterable[PatchInfo] | None = None) → Generator[PatchDetails, None, None][source]#

The context manager apply patches, usually before exporting a model.

from yobx.torch import apply_patches_for_model

with apply_patches_for_model(patch_transformers=True, model=model):
    # ...
Parameters:

  • patch_torch – applies patches for torch

  • patch_transformers – applies patch for transformers

  • verbose – prints out which patch is applies

  • model – modifies the list of patches for a particular model, it is recommended to fill it the used rope is not the default one

  • extra_patches – additional PatchInfo objects to apply alongside the built-in patches. Each patch is applied in order after the built-in ones and removed on exit.

The following shows how to use the output of this function to display information about the patches applied to the model.

from yobx.torch import apply_patches_for_model

with apply_patches_for_model(patch_transformers=True, model=model) as details:
    for patch in details:
        diff = patch.make_diff()
        print(f"-- patch {patch!r}")
        print(diff)

Custom patches can be passed directly via extra_patches without any manual do/undo bookkeeping:

from yobx.helpers.patch_helper import PatchInfo
from yobx.torch import apply_patches_for_model

my_patch = PatchInfo.make(my_fn, some_module, "fn_name", family="custom")

with apply_patches_for_model(patch_torch=True, extra_patches=[my_patch]) as details:
    ep = torch.export.export(model, ...)
yobx.torch.patch.retrieve_stacktrace()[source]#

Retrieves and prints the current stack trace, avoids every torch file.