ColossalAI/colossalai/fx/tracer/meta_patch/patched_function/activation_function.py

import torch

from ...registry import meta_patched_function


@meta_patched_function.register(torch.nn.functional.relu)
def torch_nn_func_relu(input, inplace=False):
    return torch.empty(input.shape, device='meta')
[fx] refactored the file structure of patched function and module (#1238) * [fx] refactored the file structure of patched function and module * polish code 2022-07-12 07:01:58 +00:00			`import torch`
[fx] support module with bias addition (#1780) * [autoparallel] refactor tracer to fix bias addition issue * [fx] support module with bias addition * create bias_addition_module * refactor file structure * polish code * fix unit test 2022-11-01 14:53:51 +00:00
			`from ...registry import meta_patched_function`
[fx] refactored the file structure of patched function and module (#1238) * [fx] refactored the file structure of patched function and module * polish code 2022-07-12 07:01:58 +00:00

			`@meta_patched_function.register(torch.nn.functional.relu)`
			`def torch_nn_func_relu(input, inplace=False):`
[fx] support module with bias addition (#1780) * [autoparallel] refactor tracer to fix bias addition issue * [fx] support module with bias addition * create bias_addition_module * refactor file structure * polish code * fix unit test 2022-11-01 14:53:51 +00:00			`return torch.empty(input.shape, device='meta')`