def gen_py_torch_functions(out, declarations, template_path): PY_TORCH_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_torch_functions.cpp') PY_TORCH_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_torch_functions_dispatch.h') def should_bind(declaration): return (should_generate_python_binding(declaration) and declaration['mode'] != 'NN' and 'namespace' in declaration['method_of']) py_torch_functions = group_declarations_by_name(declarations, should_bind) env = create_python_bindings(py_torch_functions, has_self=False) write(out, 'python_torch_functions.cpp', PY_TORCH_FUNCTIONS_CPP, env) write(out, 'python_torch_functions_dispatch.h', PY_TORCH_DISPATCH_H, env)
def gen_py_variable_methods(out, declarations, template_path): PY_VARIABLE_METHODS_CPP = CodeTemplate.from_file(template_path + '/python_variable_methods.cpp') PY_VARIABLE_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_variable_methods_dispatch.h') def should_bind(declaration): return (should_generate_python_binding(declaration) and declaration['mode'] != 'NN' and 'Tensor' in declaration['method_of']) py_variable_methods = group_declarations_by_name(declarations, should_bind) env = create_python_bindings(py_variable_methods, True) write(out, 'python_variable_methods.cpp', PY_VARIABLE_METHODS_CPP, env) write(out, 'python_variable_methods_dispatch.h', PY_VARIABLE_DISPATCH_H, env)
def gen_py_nn_functions(out, declarations, template_path): PY_NN_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_nn_functions.cpp') PY_NN_FUNCTIONS_H = CodeTemplate.from_file(template_path + '/python_nn_functions.h') PY_NN_DISPATCH_H = CodeTemplate.from_file(template_path + '/python_nn_functions_dispatch.h') def should_bind(declaration): return (should_generate_python_binding(declaration) and declaration['mode'] == 'NN') py_nn_functions = group_declarations_by_name(declarations, should_bind) env = create_python_bindings(py_nn_functions, has_self=False, is_module=True) write(out, 'python_nn_functions.cpp', PY_NN_FUNCTIONS_CPP, env) write(out, 'python_nn_functions.h', PY_NN_FUNCTIONS_H, env) write(out, 'python_nn_functions_dispatch.h', PY_NN_DISPATCH_H, env)
def gen_py_nn_functions(out, declarations, template_path): PY_NN_FUNCTIONS_CPP = CodeTemplate.from_file(template_path + '/python_nn_functions.cpp') PY_NN_FUNCTIONS_H = CodeTemplate.from_file(template_path + '/python_nn_functions.h') PY_NN_DISPATCH_H = CodeTemplate.from_file( template_path + '/python_nn_functions_dispatch.h') py_nn_functions = get_py_nn_functions(declarations) env = create_python_bindings(py_nn_functions, has_self=False, is_module=True) write(out, 'python_nn_functions.cpp', PY_NN_FUNCTIONS_CPP, env) write(out, 'python_nn_functions.h', PY_NN_FUNCTIONS_H, env) write(out, 'python_nn_functions_dispatch.h', PY_NN_DISPATCH_H, env)
'sub(Tensor, Scalar, Scalar)', 'sub_(Tensor, Scalar, Scalar)', 'mul(Tensor, Scalar)', 'mul_(Tensor, Scalar)', 'div(Tensor, Scalar)', 'div_(Tensor, Scalar)', ] PY_VARIABLE_METHOD_VARARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { HANDLE_TH_ERRORS static PythonArgParser parser({ ${signatures} }, /*traceable=*/${traceable}); ${unpack_self} ParsedArgs<${max_args}> parsed_args; auto r = parser.parse(args, kwargs, parsed_args); ${dispatch} Py_RETURN_NONE; END_HANDLE_TH_ERRORS } """) PY_VARIABLE_METHOD_NOARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args) { HANDLE_TH_ERRORS ${unpack_self} return wrap(${dispatch_name}(${actuals})); END_HANDLE_TH_ERRORS
'__xor__', '__rxor__', '__ixor__', '__or__', '__ror__', '__ior__', ] PY_VARIABLE_METHOD_VARARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) { HANDLE_TH_ERRORS static PythonArgParser parser({ ${signatures} }, /*traceable=*/${traceable}); ${unpack_self} ParsedArgs<${max_args}> parsed_args; auto r = parser.parse(args, kwargs, parsed_args); ${declare_namedtuple_return_types} ${dispatch} Py_RETURN_NONE; END_HANDLE_TH_ERRORS } """) PY_VARIABLE_METHOD_NOARGS = CodeTemplate("""\ static PyObject * ${pycname}(PyObject* self_, PyObject* args) { HANDLE_TH_ERRORS ${declare_namedtuple_return_types} ${unpack_self} return wrap(${namedtuple_return_type}${dispatch_name}(${actuals}));
] 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 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) def split_name_params(prototype): name, params = re.match('(\w+)\((.*)\)', prototype).groups() return name, params.split(', ') # When tracing, we record inplace operations as out-of-place operations, # because we don't have a story for side effects in the IR yet.
CodeTemplate = import_module('code_template', 'aten/src/ATen/code_template.py').CodeTemplate # You should use these lines, rather than doing it manually. # Especially if you see this error! # # File "/usr/local/lib/python2.7/dist-packages/yaml/__init__.py", line 69, in load # loader = Loader(stream) # TypeError: 'module' object is not callable 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 ${filename}") # Matches "foo" in "foo, bar" but not "foobar". Used to search for the # occurrence 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) def split_name_params(prototype): name, overload_name, params = re.match(r'(\w+)(\.\w+)?\((.*)\)', prototype).groups() return name, params.split(', ') # When tracing, we record inplace operations as out-of-place operations,