import argparse import copy import os import re import yaml from collections import defaultdict from tools.shared.module_loader import import_module from .nested_dict import nested_dict CodeTemplate = import_module('code_template', 'torch/lib/ATen/code_template.py').CodeTemplate try: # use faster C loader if available from yaml import CLoader as Loader except ImportError: from yaml import Loader METHOD_DECLARATION = CodeTemplate("""\ virtual ${return_type} ${method_prefix}${api_name}(${formals}) const override; """) METHOD_DEFINITION = CodeTemplate("""\ ${return_type} VariableType::${method_prefix}${api_name}(${formals}) const { ${type_definition_body} } """) METHOD_DEFINITION_NYI = CodeTemplate("""\ throw std::runtime_error("${api_name}: NYI");""")
# Generates Python bindings for ATen functions # # The bindings are generated as methods on python_variable or functions on the # torch._C._nn object. # from collections import defaultdict import re from .nested_dict import nested_dict from .gen_variable_type import should_trace from .utils import write try: from src.ATen.code_template import CodeTemplate except ImportError: from tools.shared.module_loader import import_module CodeTemplate = import_module('code_template', 'aten/src/ATen/code_template.py').CodeTemplate # These functions require manual Python bindings or are not exposed to Python SKIP_PYTHON_BINDINGS = [ 'alias', 'contiguous', 'is_cuda', 'is_sparse', 'size', 'stride', '.*_backward', '.*_backward_(out|input|weight|bias)', '.*_forward', '.*_forward_out', 'sparse_raw_resize_', '_unsafe_view', 'tensor', 'sparse_coo_tensor', 'th_sparse_coo_tensor', 'native_sparse_coo_tensor', '_arange.*', '_range.*', '_linspace.*', '_logspace.*', '_sparse_add.*', '_sparse_div.*', '_sparse_mul.*', '_sparse_sub.*', 'index', '_indexCopy_', 'max_values', 'min_values', 'argmax', 'argmin', '_cumsum.*', '_cumprod.*', '_sum.*', '_prod.*', '_th_.*', 'arange.*', 'range.*', '_gesv.*', '_getri.*', 'slice', '_local_scalar', '_local_scalar_dense', 'max_pool1d', 'max_pool2d', 'max_pool3d' ]
import re import os from tools.shared.module_loader import import_module from .nested_dict import nested_dict __all__ = [ 'CodeTemplate', 'IDENT_REGEX', 'YamlLoader', 'nested_dict', 'split_name_params', 'write', ] CodeTemplate = import_module('code_template', 'aten/src/ATen/code_template.py').CodeTemplate try: # use faster C loader if available from yaml import CLoader as YamlLoader except ImportError: from yaml import Loader as YamlLoader GENERATED_COMMENT = CodeTemplate( "@" + "generated from tools/autograd/templates/${filename}") # Matches "foo" in "foo, bar" but not "foobar". Used to search for the # occurence of a parameter in the derivative formula IDENT_REGEX = r'(^|\W){}($|\W)' # TODO: Use a real parser here; this will get bamboozled # by signatures that contain things like std::array<bool, 2> (note the space)