mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
d1c092ae1b
Summary: We've made the following changes: - The new way to use the API is `m.impl_abstract_pystub(module, context)`. Every subsequent m.def of an op inside the TORCH_LIBRARY block gives the op the `impl_abstract_pystub`. - Added a mechanism to determine if an operator was defined in Python or C++. Library.define in Python appends the op to a global set, which is analogous to what we do for tracking Library.impl. - If someone does `torch.library.impl_abstract` in Python for an operator, then we require that it has an `impl_abstract_pystub` specified and we also check that the module in the `impl_abstract_pystub` is the same as the module where the call to `torch.library.impl_abstract` exists. - Unfortunately we can't check the "context" (which is the buck target on buck-based systems) because buck sits above us. bypass-github-export-checks Test Plan: - existing tests Differential Revision: D51080493 Pull Request resolved: https://github.com/pytorch/pytorch/pull/113182 Approved by: https://github.com/ezyang
18 lines
460 B
Python
18 lines
460 B
Python
import torch
|
|
from model import get_custom_op_library_path
|
|
|
|
torch.ops.load_library(get_custom_op_library_path())
|
|
|
|
|
|
# NB: The impl_abstract_pystub for cos actually
|
|
# specifies it should live in the my_custom_ops2 module.
|
|
@torch.library.impl_abstract("custom::cos")
|
|
def cos_abstract(x):
|
|
return torch.empty_like(x)
|
|
|
|
|
|
# NB: There is no impl_abstract_pystub for tan
|
|
@torch.library.impl_abstract("custom::tan")
|
|
def tan_abstract(x):
|
|
return torch.empty_like(x)
|