| | import functools |
| | from inspect import signature |
| | from .common_op_utils import _basic_validation |
| |
|
| | """ |
| | Common utilities to register ops on ShardedTensor, ReplicatedTensor |
| | and PartialTensor. |
| | """ |
| |
|
| | def _register_op(op, func, op_table): |
| | """ |
| | Performs basic validation and registers the provided op in the given |
| | op_table. |
| | """ |
| | if len(signature(func).parameters) != 4: |
| | raise TypeError( |
| | f'Custom sharded op function expects signature: ' |
| | f'(types, args, kwargs, process_group), but received ' |
| | f'signature: {signature(func)}') |
| |
|
| | op_table[op] = func |
| |
|
| | def _decorator_func(wrapped_func, op, op_table): |
| | """ |
| | Decorator function to register the given ``op`` in the provided |
| | ``op_table`` |
| | """ |
| |
|
| | @functools.wraps(wrapped_func) |
| | def wrapper(types, args, kwargs, process_group): |
| | _basic_validation(op, args, kwargs) |
| | return wrapped_func(types, args, kwargs, process_group) |
| |
|
| | _register_op(op, wrapper, op_table) |
| | return wrapper |
| |
|
