xlnstorch.implements

Contents

xlnstorch.implements#

xlnstorch.implements(torch_function, lns_operation, key=None, default=False)#

A decorator to register a custom implementation for a given torch function.

This allows functions to be mapped to specific handlers in the LNS context and optionally set as the default implementation for that function.

Parameters:

  • torch_function (Callable) – The torch function to be overriden.

  • lns_operation (Callable) – The function that implements the given LNS operation to tensors. This should be the computation that is performed on the internal torch tensor representations for the LNSTensor objects.

  • key (str, optional) – A unique key to identify the implementation. If not provided, the function’s name will be used by default.

  • default (bool, optional) – If True, this implementation will be set as the default for the specified torch function. Defaults to False.

Returns:

The decorator that registers the function as an implementation for the specific torch function.

Return type:

Callable

Contents