xlnstorch.override_implementation

xlnstorch.override_implementation#

xlnstorch.override_implementation(torch_function, impl_key)#

Temporarily override the default implementation for a torch function within a context. This allows for testing or using a different implementation without permanently changing the default.

Parameters:

  • torch_function (Callable) – The torch function for which the implementation is to be temporarily overridden.

  • impl_key (str) – The key identifying the new implementation to use as default.

Yields:

None – The function yields control back to the context block.

Return type:

Generator[None, None, None]

Examples

>>> with override_implementation(torch.add, 'custom_add_impl'):
>>>     # Inside this block, torch.add will use 'custom_add_impl'
>>>     pass
Contents