mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
Run Black on all of tools/
Signed-off-by: Edward Z. Yang <ezyangfb.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/76089 Approved by: https://github.com/albanD
This commit is contained in:
Edward Z. Yang
committed by
PyTorch MergeBot
PyTorch MergeBot
parent
ae864d4fb9
commit
a11c1bbdd0
@ -11,15 +11,15 @@ try:
|
||||
except ImportError:
|
||||
from yaml import SafeLoader as YamlLoader # type: ignore[misc]
|
||||
|
||||
source_files = {'.py', '.cpp', '.h'}
|
||||
source_files = {".py", ".cpp", ".h"}
|
||||
|
||||
NATIVE_FUNCTIONS_PATH = 'aten/src/ATen/native/native_functions.yaml'
|
||||
NATIVE_FUNCTIONS_PATH = "aten/src/ATen/native/native_functions.yaml"
|
||||
|
||||
# TODO: This is a little inaccurate, because it will also pick
|
||||
# up setup_helper scripts which don't affect code generation
|
||||
def all_generator_source() -> List[str]:
|
||||
r = []
|
||||
for directory, _, filenames in os.walk('tools'):
|
||||
for directory, _, filenames in os.walk("tools"):
|
||||
for f in filenames:
|
||||
if os.path.splitext(f)[1] in source_files:
|
||||
full = os.path.join(directory, f)
|
||||
@ -27,25 +27,26 @@ def all_generator_source() -> List[str]:
|
||||
return sorted(r)
|
||||
|
||||
|
||||
def generate_code(ninja_global: Optional[str] = None,
|
||||
native_functions_path: Optional[str] = None,
|
||||
install_dir: Optional[str] = None,
|
||||
subset: Optional[str] = None,
|
||||
disable_autograd: bool = False,
|
||||
force_schema_registration: bool = False,
|
||||
operator_selector: Any = None) -> None:
|
||||
def generate_code(
|
||||
ninja_global: Optional[str] = None,
|
||||
native_functions_path: Optional[str] = None,
|
||||
install_dir: Optional[str] = None,
|
||||
subset: Optional[str] = None,
|
||||
disable_autograd: bool = False,
|
||||
force_schema_registration: bool = False,
|
||||
operator_selector: Any = None,
|
||||
) -> None:
|
||||
from tools.autograd.gen_autograd import gen_autograd, gen_autograd_python
|
||||
from tools.autograd.gen_annotated_fn_args import gen_annotated
|
||||
from tools.codegen.selective_build.selector import SelectiveBuilder
|
||||
|
||||
|
||||
# Build ATen based Variable classes
|
||||
if install_dir is None:
|
||||
install_dir = 'torch/csrc'
|
||||
python_install_dir = 'torch/testing/_internal/generated'
|
||||
install_dir = "torch/csrc"
|
||||
python_install_dir = "torch/testing/_internal/generated"
|
||||
else:
|
||||
python_install_dir = install_dir
|
||||
autograd_gen_dir = os.path.join(install_dir, 'autograd', 'generated')
|
||||
autograd_gen_dir = os.path.join(install_dir, "autograd", "generated")
|
||||
for d in (autograd_gen_dir, python_install_dir):
|
||||
os.makedirs(d, exist_ok=True)
|
||||
autograd_dir = os.fspath(pathlib.Path(__file__).parent.parent / "autograd")
|
||||
@ -54,7 +55,8 @@ def generate_code(ninja_global: Optional[str] = None,
|
||||
gen_autograd_python(
|
||||
native_functions_path or NATIVE_FUNCTIONS_PATH,
|
||||
autograd_gen_dir,
|
||||
autograd_dir)
|
||||
autograd_dir,
|
||||
)
|
||||
|
||||
if operator_selector is None:
|
||||
operator_selector = SelectiveBuilder.get_nop_selector()
|
||||
@ -73,17 +75,18 @@ def generate_code(ninja_global: Optional[str] = None,
|
||||
gen_annotated(
|
||||
native_functions_path or NATIVE_FUNCTIONS_PATH,
|
||||
python_install_dir,
|
||||
autograd_dir)
|
||||
autograd_dir,
|
||||
)
|
||||
|
||||
|
||||
def get_selector_from_legacy_operator_selection_list(
|
||||
selected_op_list_path: str,
|
||||
selected_op_list_path: str,
|
||||
) -> Any:
|
||||
with open(selected_op_list_path, 'r') as f:
|
||||
with open(selected_op_list_path, "r") as f:
|
||||
# strip out the overload part
|
||||
# It's only for legacy config - do NOT copy this code!
|
||||
selected_op_list = {
|
||||
opname.split('.', 1)[0] for opname in yaml.load(f, Loader=YamlLoader)
|
||||
opname.split(".", 1)[0] for opname in yaml.load(f, Loader=YamlLoader)
|
||||
}
|
||||
|
||||
# Internal build doesn't use this flag any more. Only used by OSS
|
||||
@ -96,6 +99,7 @@ def get_selector_from_legacy_operator_selection_list(
|
||||
is_used_for_training = True
|
||||
|
||||
from tools.codegen.selective_build.selector import SelectiveBuilder
|
||||
|
||||
selector = SelectiveBuilder.from_legacy_op_registration_allow_list(
|
||||
selected_op_list,
|
||||
is_root_operator,
|
||||
@ -114,10 +118,12 @@ def get_selector(
|
||||
sys.path.insert(0, root)
|
||||
from tools.codegen.selective_build.selector import SelectiveBuilder
|
||||
|
||||
assert not (selected_op_list_path is not None and
|
||||
operators_yaml_path is not None), \
|
||||
("Expected at most one of selected_op_list_path and " +
|
||||
"operators_yaml_path to be set.")
|
||||
assert not (
|
||||
selected_op_list_path is not None and operators_yaml_path is not None
|
||||
), (
|
||||
"Expected at most one of selected_op_list_path and "
|
||||
+ "operators_yaml_path to be set."
|
||||
)
|
||||
|
||||
if selected_op_list_path is None and operators_yaml_path is None:
|
||||
return SelectiveBuilder.get_nop_selector()
|
||||
@ -128,43 +134,43 @@ def get_selector(
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description='Autogenerate code')
|
||||
parser.add_argument('--native-functions-path')
|
||||
parser.add_argument('--ninja-global')
|
||||
parser.add_argument('--install_dir')
|
||||
parser = argparse.ArgumentParser(description="Autogenerate code")
|
||||
parser.add_argument("--native-functions-path")
|
||||
parser.add_argument("--ninja-global")
|
||||
parser.add_argument("--install_dir")
|
||||
parser.add_argument(
|
||||
'--subset',
|
||||
help='Subset of source files to generate. Can be "libtorch" or "pybindings". Generates both when omitted.'
|
||||
"--subset",
|
||||
help='Subset of source files to generate. Can be "libtorch" or "pybindings". Generates both when omitted.',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--disable-autograd',
|
||||
"--disable-autograd",
|
||||
default=False,
|
||||
action='store_true',
|
||||
help='It can skip generating autograd related code when the flag is set',
|
||||
action="store_true",
|
||||
help="It can skip generating autograd related code when the flag is set",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--selected-op-list-path',
|
||||
help='Path to the YAML file that contains the list of operators to include for custom build.',
|
||||
"--selected-op-list-path",
|
||||
help="Path to the YAML file that contains the list of operators to include for custom build.",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--operators_yaml_path',
|
||||
help='Path to the model YAML file that contains the list of operators to include for custom build.',
|
||||
"--operators_yaml_path",
|
||||
help="Path to the model YAML file that contains the list of operators to include for custom build.",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--force_schema_registration',
|
||||
action='store_true',
|
||||
help='force it to generate schema-only registrations for ops that are not'
|
||||
'listed on --selected-op-list'
|
||||
"--force_schema_registration",
|
||||
action="store_true",
|
||||
help="force it to generate schema-only registrations for ops that are not"
|
||||
"listed on --selected-op-list",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gen_lazy_ts_backend',
|
||||
action='store_true',
|
||||
help='Enable generation of the torch::lazy TorchScript backend'
|
||||
"--gen_lazy_ts_backend",
|
||||
action="store_true",
|
||||
help="Enable generation of the torch::lazy TorchScript backend",
|
||||
)
|
||||
parser.add_argument(
|
||||
'--per_operator_headers',
|
||||
action='store_true',
|
||||
help='Build lazy tensor ts backend with per-operator ATen headers, must match how ATen was built'
|
||||
"--per_operator_headers",
|
||||
action="store_true",
|
||||
help="Build lazy tensor ts backend with per-operator ATen headers, must match how ATen was built",
|
||||
)
|
||||
options = parser.parse_args()
|
||||
|
||||
@ -176,12 +182,14 @@ def main() -> None:
|
||||
options.disable_autograd,
|
||||
options.force_schema_registration,
|
||||
# options.selected_op_list
|
||||
operator_selector=get_selector(options.selected_op_list_path, options.operators_yaml_path),
|
||||
operator_selector=get_selector(
|
||||
options.selected_op_list_path, options.operators_yaml_path
|
||||
),
|
||||
)
|
||||
|
||||
if options.gen_lazy_ts_backend:
|
||||
aten_path = os.path.dirname(os.path.dirname(options.native_functions_path))
|
||||
ts_backend_yaml = os.path.join(aten_path, 'native/ts_native_functions.yaml')
|
||||
ts_backend_yaml = os.path.join(aten_path, "native/ts_native_functions.yaml")
|
||||
ts_native_functions = "torch/csrc/lazy/ts_backend/ts_native_functions.cpp"
|
||||
ts_node_base = "torch/csrc/lazy/ts_backend/ts_node.h"
|
||||
if options.install_dir is None:
|
||||
@ -189,22 +197,29 @@ def main() -> None:
|
||||
lazy_install_dir = os.path.join(options.install_dir, "lazy/generated")
|
||||
os.makedirs(lazy_install_dir, exist_ok=True)
|
||||
|
||||
assert os.path.isfile(ts_backend_yaml), f"Unable to access ts_backend_yaml: {ts_backend_yaml}"
|
||||
assert os.path.isfile(ts_native_functions), f"Unable to access {ts_native_functions}"
|
||||
assert os.path.isfile(
|
||||
ts_backend_yaml
|
||||
), f"Unable to access ts_backend_yaml: {ts_backend_yaml}"
|
||||
assert os.path.isfile(
|
||||
ts_native_functions
|
||||
), f"Unable to access {ts_native_functions}"
|
||||
from tools.codegen.gen_lazy_tensor import run_gen_lazy_tensor
|
||||
from tools.codegen.dest.lazy_ir import TSLazyIR
|
||||
run_gen_lazy_tensor(aten_path=aten_path,
|
||||
source_yaml=ts_backend_yaml,
|
||||
backend_name="TorchScript",
|
||||
output_dir=lazy_install_dir,
|
||||
dry_run=False,
|
||||
impl_path=ts_native_functions,
|
||||
node_base="TsNode",
|
||||
node_base_hdr=ts_node_base,
|
||||
build_in_tree=True,
|
||||
lazy_ir_cls=TSLazyIR,
|
||||
per_operator_headers=options.per_operator_headers,
|
||||
gen_forced_fallback_code=True)
|
||||
|
||||
run_gen_lazy_tensor(
|
||||
aten_path=aten_path,
|
||||
source_yaml=ts_backend_yaml,
|
||||
backend_name="TorchScript",
|
||||
output_dir=lazy_install_dir,
|
||||
dry_run=False,
|
||||
impl_path=ts_native_functions,
|
||||
node_base="TsNode",
|
||||
node_base_hdr=ts_node_base,
|
||||
build_in_tree=True,
|
||||
lazy_ir_cls=TSLazyIR,
|
||||
per_operator_headers=options.per_operator_headers,
|
||||
gen_forced_fallback_code=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user