utils.py 54.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

17
import ast
18
import astor
19 20
import atexit
import copy
21
import collections
22
from paddle.utils import gast
23 24 25 26
import inspect
import os
import six
import tempfile
27
import textwrap
28
import numpy as np
29

30
import paddle
31
from paddle.fluid import unique_name
32
from paddle.fluid.data_feeder import convert_dtype
33

34 35 36 37 38 39
# Note(Aurelius): Do not forget the dot `.` to distinguish other
# module such as paddlenlp.
PADDLE_MODULE_PREFIX = 'paddle.'
DYGRAPH_MODULE_PREFIX = 'paddle.fluid.dygraph'
DYGRAPH_TO_STATIC_MODULE_PREFIX = 'paddle.fluid.dygraph.dygraph_to_static'

40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61

class BaseNodeVisitor(gast.NodeVisitor):
    """
    Implement customized NodeVisitor inherited from gast.NodeVisitor. 
    Ancestor nodes are traced to easily support more operations of currently
    visited node.
    """

    def __init__(self):
        self.ancestor_nodes = []

    def visit(self, node):
        """Visit a node."""
        self.ancestor_nodes.append(node)

        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
        ret = visitor(node)
        self.ancestor_nodes.pop()
        return ret


62
# imp is deprecated in python3
T
tianshuo78520a 已提交
63
from importlib.machinery import SourceFileLoader
64

65 66 67 68 69 70 71 72 73 74
dygraph_class_to_static_api = {
    "CosineDecay": "cosine_decay",
    "ExponentialDecay": "exponential_decay",
    "InverseTimeDecay": "inverse_time_decay",
    "NaturalExpDecay": "natural_exp_decay",
    "NoamDecay": "noam_decay",
    "PiecewiseDecay": "piecewise_decay",
    "PolynomialDecay": "polynomial_decay",
}

75
FOR_ITER_INDEX_PREFIX = '__for_loop_var_index'
76 77
FOR_ITER_TUPLE_PREFIX = '__for_loop_iter_tuple'
FOR_ITER_TUPLE_INDEX_PREFIX = '__for_loop_iter_tuple_index'
78
FOR_ITER_VAR_LEN_PREFIX = '__for_loop_var_len'
79
FOR_ITER_VAR_NAME_PREFIX = '__for_loop_iter_var'
80
FOR_ITER_ZIP_TO_LIST_PREFIX = '__for_loop_iter_zip'
81

82 83 84 85 86 87 88 89
# FullArgSpec is valid from Python3. Defined a Namedtuple to
# to make it available in Python2.
FullArgSpec = collections.namedtuple('FullArgSpec', [
    'args', 'varargs', 'varkw', 'defaults', 'kwonlyargs', 'kwonlydefaults',
    'annotations'
])


90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
class UndefinedVar:

    def __init__(self, name):
        self.name = name

    def check(self):
        raise UnboundLocalError(
            "local variable '{}' should be created before using it.")


def saw(x):
    if isinstance(x, UndefinedVar):
        return x.check()
    else:
        return x


107 108 109 110 111
def getfullargspec(target):
    if hasattr(inspect, "getfullargspec"):
        return inspect.getfullargspec(target)
    else:
        argspec = inspect.getargspec(target)
112 113 114 115 116 117 118
        return FullArgSpec(args=argspec.args,
                           varargs=argspec.varargs,
                           varkw=argspec.keywords,
                           defaults=argspec.defaults,
                           kwonlyargs=[],
                           kwonlydefaults=None,
                           annotations={})
119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140


def parse_arg_and_kwargs(function):
    """
    Returns full argument names as list. e.g ['x', 'y', 'z']
    """
    fullargspec = getfullargspec(function)
    arg_names = fullargspec.args
    if arg_names and 'self' == arg_names[0]:
        arg_names = fullargspec.args[1:]

    # parse default kwargs
    default_kwargs = {}
    default_values = fullargspec.defaults
    if default_values:
        assert len(default_values) <= len(arg_names)
        default_kwarg_names = arg_names[-len(default_values):]
        default_kwargs = dict(zip(default_kwarg_names, default_values))

    return arg_names, default_kwargs


W
WeiXin 已提交
141 142 143 144 145 146 147 148 149
def parse_varargs_name(function):
    """
    Returns varargs name string of function. e.g: 'input' from `foo(x, *input)`
    """
    fullargspec = getfullargspec(function)
    varargs = fullargspec.varargs
    return varargs


150 151 152 153 154 155 156 157
def type_name(v):
    return type(v).__name__


def make_hashable(x, error_msg=None):
    """
    Makes input `x` hashable.

158
    For some unhashable objects, such as `dict/list/set/np.ndarray`,applying hash function by using their values.
159
    """
160
    if isinstance(x, (tuple, list, set)):
161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177
        return tuple(map(make_hashable, x))

    try:
        hash(x)
    except TypeError:
        if isinstance(x, np.ndarray):
            # Note: `tostring()` will return the binary data from np.ndarray that
            # means different value will lead to different hash code.
            return hash(x.tostring())
        elif isinstance(x, dict):
            return tuple(map(make_hashable, x.values()))

        error_msg = error_msg or "Requires a hashable object."
        raise ValueError(error_msg + " But received type: %s" % type_name(x))

    return x

178

179 180 181 182 183 184 185
def _is_api_in_module_helper(obj, module_prefix):
    m = inspect.getmodule(obj)
    return m is not None and m.__name__.startswith(module_prefix)


def is_api_in_module(node, module_prefix):
    assert isinstance(node, gast.Call), "Input non-Call node for is_dygraph_api"
186 187 188 189 190 191 192 193

    # Python can have gast.Call as function, for example: covert_call(func)(x)
    # We only check the most outside function
    func_node = node.func
    while isinstance(func_node, gast.Call):
        func_node = func_node.func

    func_str = astor.to_source(gast.gast_to_ast(func_node)).strip()
194
    try:
195 196 197 198 199
        # TODO(liym27):
        #  Consider a better to import modules like:
        #  source_file = inspect.getfile(dyfunc)
        #  import_statements = ImportVisitor(source_file).transform()
        #  import_str = "".join(import_statements)
200
        import paddle
L
liym27 已提交
201
        import paddle.fluid as fluid
202
        import paddle.fluid.dygraph as dygraph
L
liym27 已提交
203
        import paddle.fluid.layers as layers
204
        import paddle.jit.dy2static as _jst
205

206
        from paddle.fluid.dygraph import to_variable
207 208
        from paddle import to_tensor

209 210
        return eval("_is_api_in_module_helper({}, '{}')".format(
            func_str, module_prefix))
211
    except Exception:
212 213 214 215
        return False


def is_dygraph_api(node):
216

217
    # Note: A api in module dygraph_to_static is not a real dygraph api.
218
    if is_api_in_module(node, DYGRAPH_TO_STATIC_MODULE_PREFIX):
219 220
        return False

221 222
    # TODO(liym27): A better way to determine whether it is a dygraph api.
    #  Consider the decorator @dygraph_only
223
    return is_api_in_module(node, DYGRAPH_MODULE_PREFIX)
224 225 226


def is_paddle_api(node):
227 228 229 230 231 232
    return is_api_in_module(node, PADDLE_MODULE_PREFIX)


def is_paddle_func(func):
    m = inspect.getmodule(func)
    return m is not None and m.__name__.startswith(PADDLE_MODULE_PREFIX)
233 234 235 236 237 238 239 240 241 242 243 244 245 246


# Is numpy_api cannot reuse is_api_in_module because of numpy module problem
def is_numpy_api(node):
    assert isinstance(node, gast.Call), "Input non-Call node for is_numpy_api"
    func_str = astor.to_source(gast.gast_to_ast(node.func))
    try:
        import numpy as np
        module_result = eval("_is_api_in_module_helper({}, '{}')".format(
            func_str, "numpy"))
        # BUG: np.random.uniform doesn't have module and cannot be analyzed
        # TODO: find a better way
        if not module_result:
            return func_str.startswith("numpy.") or func_str.startswith("np.")
247
    except Exception:
248 249 250
        return False


L
liym27 已提交
251 252 253
def is_control_flow_to_transform(node,
                                 static_analysis_visitor=None,
                                 var_name_to_type=None):
254
    """
L
liym27 已提交
255 256
    Determines whether the node is a PaddlePaddle control flow statement which needs to
    be transformed into a static graph control flow statement.
257 258 259
    """
    assert isinstance(node, gast.AST), \
        "The type of input node must be gast.AST, but received %s." % type(node)
260 261 262
    visitor = IsControlFlowVisitor(node,
                                   static_analysis_visitor,
                                   node_var_type_map=var_name_to_type)
L
liym27 已提交
263 264
    need_to_transform = visitor.transform()
    return need_to_transform
265 266


267 268
def _delete_keywords_from(node):
    assert isinstance(node, gast.Call)
269
    func_src = astor.to_source(gast.gast_to_ast(node.func))
270 271 272 273 274 275 276 277 278 279 280 281
    import paddle.fluid as fluid
    full_args = eval("inspect.getargspec({})".format(func_src))
    full_args_name = full_args[0]

    node.keywords = [k for k in node.keywords if k.arg in full_args_name]
    return


def to_static_api(dygraph_class):
    if dygraph_class in dygraph_class_to_static_api:
        return dygraph_class_to_static_api[dygraph_class]
    else:
282 283 284
        raise NotImplementedError(
            "Paddle dygraph API {} cannot be converted "
            "to static graph at present.".format(dygraph_class))
285 286 287 288 289 290 291 292 293 294


def _add_keywords_to(node, dygraph_api_name):
    assert isinstance(node, gast.Call)
    if dygraph_api_name == "Linear":
        for ast_keyword in node.keywords:
            if ast_keyword.arg == "output_dim":
                ast_keyword.arg = "size"

        node.keywords.append(
295 296
            gast.keyword(arg="num_flatten_dims",
                         value=gast.Constant(value=-1, kind=None)))
297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314

    if dygraph_api_name == "BilinearTensorProduct":
        for ast_keyword in node.keywords:
            if ast_keyword.arg == "output_dim":
                ast_keyword.arg = "size"

    if dygraph_api_name == "PRelu":
        for ast_keyword in node.keywords:
            if ast_keyword.arg == "input":
                ast_keyword.arg = "x"
    return


def to_static_ast(node, class_node):
    assert isinstance(node, gast.Call)
    assert isinstance(class_node, gast.Call)
    static_api = to_static_api(class_node.func.attr)

315 316 317 318 319 320 321 322 323
    node.func = gast.Attribute(attr=static_api,
                               ctx=gast.Load(),
                               value=gast.Attribute(attr='layers',
                                                    ctx=gast.Load(),
                                                    value=gast.Name(
                                                        ctx=gast.Load(),
                                                        id='fluid',
                                                        annotation=None,
                                                        type_comment=None)))
324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343

    update_args_of_func(node, class_node, 'forward')

    node.args.extend(class_node.args)
    node.keywords.extend(class_node.keywords)
    _add_keywords_to(node, class_node.func.attr)
    _delete_keywords_from(node)

    gast.fix_missing_locations(node)

    return node


def update_args_of_func(node, dygraph_node, method_name):
    assert isinstance(node, gast.Call)
    if method_name not in ["__init__", "forward"]:
        raise ValueError(
            "The method name of class to update args should be '__init__' or 'forward'"
        )

344
    class_src = astor.to_source(gast.gast_to_ast(dygraph_node.func))
345 346 347
    import paddle.fluid as fluid
    if method_name == "__init__" or eval(
            "issubclass({}, fluid.dygraph.Layer)".format(class_src)):
348 349
        full_args = eval("inspect.getargspec({}.{})".format(
            class_src, method_name))
350 351 352 353 354 355 356 357 358 359 360
        full_args_name = [
            arg_name for arg_name in full_args[0] if arg_name != "self"
        ]
    else:
        full_args_name = []
    added_keywords = []
    for idx, arg in enumerate(node.args):
        added_keywords.append(gast.keyword(arg=full_args_name[idx], value=arg))

    node.args = []
    node.keywords = added_keywords + node.keywords
361 362 363


def create_api_shape_node(tensor_shape_node):
364 365 366 367 368
    assert isinstance(tensor_shape_node,
                      (gast.Name, gast.Attribute, gast.Subscript))

    if isinstance(tensor_shape_node, gast.Name):
        api_shape_node = gast.Call(
369
            func=gast.parse('paddle.shape').body[0].value,
370 371 372
            args=[tensor_shape_node],
            keywords=[])
        return api_shape_node
373 374 375

    if isinstance(tensor_shape_node, gast.Attribute):
        api_shape_node = gast.Call(
376
            func=gast.parse('paddle.shape').body[0].value,
377 378 379 380 381 382 383 384
            args=[tensor_shape_node.value],
            keywords=[])
        return api_shape_node

    if isinstance(tensor_shape_node, gast.Subscript):
        result_node = copy.deepcopy(tensor_shape_node)
        result_node.value = create_api_shape_node(result_node.value)
        return result_node
385 386


387
def get_constant_variable_node(name, value, shape=[1], dtype='int64'):
388 389
    return gast.parse('%s = paddle.full(%s, "%s", %s)' %
                      (name, str(shape), str(value), dtype))
390 391 392 393 394 395 396 397 398


def get_attribute_full_name(node):
    assert isinstance(
        node,
        gast.Attribute), "Input non-Attribute node to get attribute full name"
    return astor.to_source(gast.gast_to_ast(node)).strip()


399
def generate_name_node(name_ids, ctx=gast.Load(), gen_tuple_if_single=False):
400
    """
401 402 403 404 405 406 407
    If name_ids is list or tuple or set with multiple strings, this function
    generates gast.Tuple of gast.Name.
    If the name_ids is single string or contains only 1 string, this function
    returns gast.Name if gen_tuple_if_single==False else returns gast.Tuple
    with only one gast.Name

    This function is used at several gast.Return statements.
408 409 410 411
    """
    if isinstance(name_ids, six.string_types):
        name_ids = [name_ids]
    if not isinstance(name_ids, (list, tuple, set)):
412 413 414
        raise TypeError(
            'name_ids must be list or tuple or set, but received %s' %
            type(type(name_ids)))
415
    gast_names = [
416
        gast.Name(id=name_id, ctx=ctx, annotation=None, type_comment=None)
417 418
        for name_id in name_ids
    ]
419
    if len(gast_names) == 1 and not gen_tuple_if_single:
420 421 422 423 424 425 426 427 428 429 430 431 432
        name_node = gast_names[0]
    else:
        name_node = gast.Tuple(elts=gast_names, ctx=ctx)
    return name_node


def create_funcDef_node(nodes, name, input_args, return_name_ids):
    """
    Wrapper all statements of nodes into one ast.FunctionDef, which can be
    called by ast.Call.
    """
    nodes = copy.copy(nodes)
    # add return statement
433 434
    if return_name_ids:
        nodes.append(gast.Return(value=generate_name_node(return_name_ids)))
435 436
    else:
        nodes.append(gast.Return(value=None))
437 438 439 440 441 442
    func_def_node = gast.FunctionDef(name=name,
                                     args=input_args,
                                     body=nodes,
                                     decorator_list=[],
                                     returns=None,
                                     type_comment=None)
443 444 445
    return func_def_node


446 447 448 449 450 451 452 453
def index_in_list(array_list, item):
    try:
        return array_list.index(item)
    except ValueError:
        # Item not in array_list
        return -1


454 455 456 457 458 459 460 461 462 463
def create_assign_node(name, node):
    """
    Creates a `gast.Assign` node by given name_id as target and node as value.
    """
    targets = generate_name_node(name, ctx=gast.Store())
    assign_node = gast.Assign(targets=[targets], value=node)
    return targets, assign_node


class RenameTransformer(gast.NodeTransformer):
464

465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491
    def __init__(self, node):
        assert isinstance(
            node, gast.AST), "RenameTransformer only accepts gast.AST as input"
        self.root = node
        self.old_name = ""
        self.new_name = ""

    def rename(self, old_name, new_name):
        self.old_name = old_name
        self.new_name = new_name
        self.visit(self.root)

    def visit_Name(self, node):
        self.generic_visit(node)
        if node.id == self.old_name:
            node.id = self.new_name
        return node

    def visit_Attribute(self, node):
        self.generic_visit(node)
        attr_full_name = get_attribute_full_name(node)
        if attr_full_name == self.old_name:
            new_name_node = gast.parse(self.new_name).body[0].value
            return new_name_node
        return node


492
def ast_to_func(ast_root, dyfunc, delete_on_exit=True):
493 494
    """
    Transform modified AST of decorated function into python callable object.
495 496
    TODO: If only decorate one of inner function instead of decorating the main
    function, the other inner functions are invisible for the decorated function.
497
    """
498

499
    def remove_if_exit(filepath):
500 501 502
        if os.path.exists(filepath):
            os.remove(filepath)

503
    source = ast_to_source_code(ast_root)
504
    source = _inject_import_statements() + source
505

506 507 508 509
    f = tempfile.NamedTemporaryFile(mode='w',
                                    suffix='.py',
                                    delete=False,
                                    encoding='utf-8')
510 511 512 513 514
    with f:
        module_name = os.path.basename(f.name[:-3])
        f.write(source)

    if delete_on_exit:
515 516
        atexit.register(lambda: remove_if_exit(f.name))
        atexit.register(lambda: remove_if_exit(f.name[:-3] + ".pyc"))
517

T
tianshuo78520a 已提交
518
    module = SourceFileLoader(module_name, f.name).load_module()
519
    func_name = dyfunc.__name__
W
WeiXin 已提交
520 521 522 523 524 525 526 527
    # The 'forward' or 'another_forward' of 'TranslatedLayer' cannot be obtained
    # through 'func_name'. So set the special function name '__i_m_p_l__'.
    if hasattr(module, '__i_m_p_l__'):
        callable_func = getattr(module, '__i_m_p_l__')
        callable_func.__name__ = func_name
    elif hasattr(module, func_name):
        callable_func = getattr(module, func_name)
    else:
528 529 530
        raise ValueError(
            'Function: %s doesn\'t exist in the Module transformed from AST.' %
            func_name)
531 532 533 534 535 536 537 538
    # After transform dygraph function into callable_func saved in tmp file,
    # it lost the global variables from imported statements or defined in source file.
    # Recovers the necessary variables by `__globals__`.
    recover_globals_attribute(dyfunc, callable_func)

    return callable_func, f.name


539 540
def _inject_import_statements():
    import_statements = [
541
        "import paddle", "from paddle import Tensor",
542 543
        "import paddle.fluid as fluid", "import paddle.jit.dy2static as _jst",
        "from typing import *", "import numpy as np"
544 545 546 547
    ]
    return '\n'.join(import_statements) + '\n'


548 549 550 551 552
def recover_globals_attribute(src_obj, dst_obj):
    attr_name = '__globals__'

    src_globals = getattr(src_obj, attr_name, {})
    dst_globals = getattr(dst_obj, attr_name, {})
553

554
    for k, v in six.iteritems(src_globals):
555 556 557
        # ignore builtin attribute.
        if not (k.startswith('__') and k.endswith('__')):
            dst_globals[k] = v
558 559


560 561 562 563 564 565
def func_to_source_code(function, dedent=True):
    """
    Transforms function into raw string of source code.
    """
    if not (inspect.isfunction(function) or inspect.ismethod(function)):
        raise TypeError(
566 567
            "The type of 'function' should be a function or method, but received {}."
            .format(type(function).__name__))
568
    source_code_list, _ = inspect.getsourcelines(function)
569
    # Replace comments with blank lines so that error messages are not misplaced
570
    source_code_list = [
571 572
        line if not line.lstrip().startswith('#') else '\n'
        for line in source_code_list
573 574
    ]
    source_code = ''.join(source_code_list)
575 576 577 578 579 580
    if dedent:
        source_code = textwrap.dedent(source_code)

    return source_code


581 582
def ast_to_source_code(ast_node):
    """
583
    Transforms ast node into source code.
584 585 586 587 588 589 590 591 592
    """
    if not isinstance(ast_node, (gast.AST, ast.AST)):
        raise TypeError(
            "Type of ast_root should be gast.AST or ast.AST, but received %s." %
            type(ast_node))
    if isinstance(ast_node, gast.AST):
        ast_node = gast.gast_to_ast(ast_node)
    source_code = astor.to_source(ast_node)
    return source_code
L
liym27 已提交
593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615


def is_candidate_node(node):
    """
    Nodes with specified type will be dependent on tensor.
    """
    is_compare_node = isinstance(node, (gast.Compare, gast.BoolOp, gast.UnaryOp,
                                        gast.For, gast.If, gast.While))
    # TODO(Aurelius84): `.numpy()` may be an customized function,
    # and should consider a more elegant way to solve this problem.
    has_numpy_attr = ".numpy()" in ast_to_source_code(node)
    return is_compare_node or has_numpy_attr


def compare_with_none(node):
    """
    Whether the comparator of `gast.Compare` node is `None`.
    """
    if isinstance(node, gast.Compare):
        for child in [node.left, node.comparators]:
            # node.comparators is a list.
            if isinstance(child, list):
                child = child[0]
616 617 618
            if (isinstance(child, gast.Constant)
                    and child.value is None) or (isinstance(child, gast.Name)
                                                 and child.id == 'None'):
L
liym27 已提交
619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635
                return True
    return False


class IsControlFlowVisitor(gast.NodeVisitor):
    """
    Judge whether the ast_node of control flow from Dygraph code dependent on paddle Tensor.
    `ast_node` can be gast.If, gast.For, gast.While, gast.If.test(gast.Compare, gast.BoolOp, gast.UnaryOp).

    If returns True,
    gast.If.test must meet at least one of the following requirements:
        1. involves at least one var whose type is Tensor.
        2. the Tensor var calls `.numpy()[]` interface or Tensor.shape is [1].
        3. involves Tensor.shape[i] and the shape[i] is unknown in compile time.
    gast.While must meet at least one of the requirements 1 to 5:
        4. has `break` statement.
        5. has `continue` statement.
636
    gast.For must meet at least one of the requirements 4 to 8:
L
liym27 已提交
637
        6. calls `range` function in `for` statement and the argument of range is Tensor.
638 639
        7. calls `enumerate` function in `for` statement and the argument of enumerate is Tensor.
        8. the iterable varaible in `for` statement is Tensor.
L
liym27 已提交
640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674
        TODO: Support non-range case

    The following examples should not be considered as control_flow_if:
        1. `if Tensor_var` or `if Tensor_var is None`
        2. if Tensor.shape[i] is determined with fixed value (not -1 or None)

    Note: pred in ConditionalBlock require variable, which means all vars should be Tensor
          or transformed into Tensor, like fill_constant(shape=[1], dtype='int32', value=Tensor.shape[i]).

    TODO: 1. need to deal with `tensor.shape[i]` which need to eval the data of shape[i],
             because reshape_op may be called before this statement.
    """

    def __init__(self,
                 ast_node,
                 static_analysis_visitor=None,
                 node_var_type_map=None):
        assert isinstance(
            ast_node, gast.AST
        ), "Type of input node should be gast.AST, but received %s." % type(
            ast_node)
        self.ast_root = ast_node
        if static_analysis_visitor is None:
            from .static_analysis import StaticAnalysisVisitor
            static_analysis_visitor = StaticAnalysisVisitor(ast_node)
        self.static_analysis_visitor = static_analysis_visitor
        self.node_to_wrapper_map = self.static_analysis_visitor.get_node_to_wrapper_map(
        )
        self.node_var_type_map = node_var_type_map

        self.is_control_flow_num = 0
        self._compare_node_tenor_set = set()

    def transform(self):
        node = self.ast_root
675 676 677 678 679 680 681 682
        if isinstance(node, gast.If):
            self._visit_If(node)
        elif isinstance(node, gast.For):
            self._visit_For(node)
        elif isinstance(node, gast.While):
            self._visit_While(node)
        else:
            self.visit(node)
L
liym27 已提交
683 684 685 686 687 688 689 690 691
        return self.is_control_flow_num > 0

    def _visit_If(self, node):
        assert isinstance(node, gast.If)
        self.visit(node.test)
        return

    def _visit_For(self, node):
        assert isinstance(node, gast.For)
692 693 694 695 696 697 698 699 700 701 702 703 704 705
        if isinstance(node.iter, gast.Call):
            # for in range(var[0]|var.numpy()[0]) or for in enumerate(var|var.numpy())
            if isinstance(node.iter.func, gast.Name):
                if node.iter.func.id == "range" or node.iter.func.id == "enumerate":
                    for arg in node.iter.args:
                        self.visit(arg)
                else:
                    return
            # for in var.numpy()
            elif isinstance(node.iter.func, gast.Attribute):
                if node.iter.func.attr == 'numpy':
                    self._visit_Call(node.iter)
                else:
                    return
706 707
            else:
                return
708 709 710
        elif isinstance(node.iter, gast.Name):
            # for in var
            self.visit(node.iter)
711
        else:
L
liym27 已提交
712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749
            return

        for child_node in gast.walk(node):
            if isinstance(child_node, (gast.Continue, gast.Break)):
                self._visit_break_continue(child_node)
        return

    def _visit_While(self, node):
        assert isinstance(node, gast.While)
        test = node.test
        self.generic_visit(test)
        for child_node in gast.walk(node):
            if isinstance(child_node, (gast.Continue, gast.Break)):
                self._visit_break_continue(child_node)
        return

    def _visit_break_continue(self, node):
        assert isinstance(node, (gast.Break, gast.Continue))
        wrapper_node = self.node_to_wrapper_map.get(node)
        if not wrapper_node:
            # Transformed node is not in node_to_wrapper_map
            return

        while wrapper_node.parent:
            parent_node = wrapper_node.parent.node
            if isinstance(parent_node, (gast.For, gast.While)):
                if parent_node is self.ast_root:
                    self.is_control_flow_num += 1
                    return
                else:
                    return

            wrapper_node = wrapper_node.parent

        return

    def visit_BoolOp(self, node):
        for i, child in enumerate(node.values):
750
            self.visit(child)
L
liym27 已提交
751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799
        return node

    def visit_Compare(self, node):
        pre_control_flow_num = self.is_control_flow_num
        if not compare_with_none(node):
            self.generic_visit(node)
            for child in gast.walk(node):
                if isinstance(child, gast.Subscript):
                    self._visit_Subscript(child)
        if self.is_control_flow_num > pre_control_flow_num:
            self._compare_node_tenor_set.add(node)
        return node

    def _visit_Subscript(self, node):
        self.generic_visit(node)
        if hasattr(node, 'value') and isinstance(node.value, gast.Call):
            self._visit_Call(node.value)
        return node

    def _visit_Call(self, node):
        assert isinstance(node, gast.Call)
        if isinstance(node.func, gast.Attribute):
            attr_node = node.func
            if attr_node.attr == 'numpy':
                self.is_control_flow_num += 1

    def visit_Call(self, node):
        self._visit_Call(node)
        if is_paddle_api(node):
            self.is_control_flow_num += 1
        return node

    def visit_Name(self, node):
        if self._is_node_with_tensor(node, node.id):
            self.is_control_flow_num += 1
        return node

    def visit_Constant(self, node):
        if self._is_node_with_tensor(node, node.value):
            self.is_control_flow_num += 1
        return node

    def _is_node_with_tensor(self, node, name_id):
        from paddle.fluid.dygraph.dygraph_to_static.static_analysis import NodeVarType

        # Look up the node_var_type_map by name_id.
        if self.node_var_type_map:
            if name_id and isinstance(name_id, six.string_types):
                var_type = self.node_var_type_map.get(name_id, None)
800
                if var_type and var_type & NodeVarType.TENSOR_TYPES:
L
liym27 已提交
801 802
                    return True
        # if not found, look up the node_to_wrapper_map by node.
803
        wrapper_node = self.node_to_wrapper_map.get(node, None)
L
liym27 已提交
804
        if wrapper_node is not None:
805
            if wrapper_node.node_var_type & NodeVarType.TENSOR_TYPES:
L
liym27 已提交
806 807 808 809 810 811
                return True

        return False

    def get_compare_nodes_with_tensor(self):
        return self._compare_node_tenor_set
812 813 814 815


class NameNodeReplaceTransformer(gast.NodeTransformer):
    """
816
    This class replaces specified gast.Name node by replace_node.
817 818 819 820
    """

    def __init__(self, root_node, target_name, replace_node):
        assert isinstance(target_name, str)
821 822 823 824 825 826 827 828 829 830 831 832 833 834

        # NOTE(liym27):
        # Use gast.Name to replace gast.Name, otherwise, errors may occur.
        #
        # For examples:
        # If using a gast.Subscript to replace gast.Name, and the original gast.Name
        # is in the arguments of FunctionDef, an exception will be raised.
        #
        # ```
        # def func(x[i])) # x[i] can not be a argument
        #    # ...
        # ```

        assert isinstance(replace_node, gast.Name)
835 836 837 838 839 840 841 842 843 844 845
        self.target_name = target_name
        self.replace_node = replace_node

        self.visit(root_node)

    def visit_Name(self, node):
        if node.id == self.target_name:
            return self.replace_node
        return node


846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889
class ForLoopTuplePreTransformer(gast.NodeTransformer):
    """
    ForNodeVisitor parses 3 type statements (Here var is VarBase(Tensor) or python variable):
        1). for x in range(var[*]|var.numpy()[*])
        2). for x in var|var.numpy()
        3). for i, x in enumerate(var|var.numpy())

        We chose these 3 types because they are easier (x can be variable name iterating in var).
        However, users can write tuples in Python for loop, such as
        1). for var1, var2 in var|var.numpy()
        2). for t in enumerate(var|var.numpy())
        2). for i, (var1, var2, va3) in enumerate(var|var.numpy())

        To handle these case, this method will do the rewrite tuple pre-process:
        1). Non-enumerate case: for var1, var2 in var|var.numpy() will be re-written as:
          for FOR_ITER_TUPLE_PREFIX_x in var | var.numpy():
            var1 = FOR_ITER_TUPLE_PREFIX_x[0]
            var2 = FOR_ITER_TUPLE_PREFIX_x[1]
        2). Enumerate out tuple case: for t in enumerate(var|var.numpy) will be rewritten as:
          for FOR_ITER_TUPLE_INDEX_PREFIX_x, FOR_ITER_TUPLE_PREFIX_x in enumerate(var|var.numpy):
            t = (FOR_ITER_TUPLE_INDEX_PREFIX_x, FOR_ITER_TUPLE_PREFIX_x)
        3). Enumerate inner tuple case: for i, (var1, (var2, va3)) in enumerate(var|var.numpy()) will
        be re-written as:
          for i, FOR_ITER_TUPLE_PREFIX_x in var | var.numpy():
            var1 = FOR_ITER_TUPLE_PREFIX_x[0]
            var2 = FOR_ITER_TUPLE_PREFIX_x[1][0]
            var3 = FOR_ITER_TUPLE_PREFIX_x[1][1]
    """

    def __init__(self, wrapper_root):
        self.wrapper_root = wrapper_root
        self.root = wrapper_root.node

    def transform(self):
        self.visit(self.root)

    def visit_For(self, node):
        if self.is_for_enumerate_iter(node):
            if isinstance(node.target, (gast.Name, gast.Attribute)):
                # Out tuple case
                out_tuple_name = ast_to_source_code(node.target).strip()
                tuple_iter_name = unique_name.generate(
                    FOR_ITER_TUPLE_INDEX_PREFIX)
                tuple_var_name = unique_name.generate(FOR_ITER_TUPLE_PREFIX)
890 891 892 893 894 895 896 897 898 899 900
                node.target = gast.Tuple(elts=[
                    gast.Name(id=tuple_iter_name,
                              ctx=gast.Store(),
                              annotation=None,
                              type_comment=None),
                    gast.Name(id=tuple_var_name,
                              ctx=gast.Store(),
                              annotation=None,
                              type_comment=None)
                ],
                                         ctx=gast.Store())
901 902
                node.body.insert(
                    0,
903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921
                    gast.Assign(targets=[
                        gast.Name(id=out_tuple_name,
                                  ctx=gast.Store(),
                                  annotation=None,
                                  type_comment=None)
                    ],
                                value=gast.Tuple(elts=[
                                    gast.Name(id=tuple_iter_name,
                                              ctx=gast.Load(),
                                              annotation=None,
                                              type_comment=None),
                                    gast.Name(id=tuple_var_name,
                                              ctx=gast.Load(),
                                              annotation=None,
                                              type_comment=None)
                                ],
                                                 ctx=gast.Load())))
            elif isinstance(node.target, (gast.List, gast.Tuple)) and len(
                    node.target.elts) >= 2 and isinstance(
922 923 924 925
                        node.target.elts[1], (gast.List, gast.Tuple)):
                # Inner tuple case
                inner_tuple_name = unique_name.generate(FOR_ITER_TUPLE_PREFIX)
                origin_inner_tuple_node = node.target.elts[1]
926 927 928 929
                node.target.elts[1] = gast.Name(id=inner_tuple_name,
                                                ctx=gast.Store(),
                                                annotation=None,
                                                type_comment=None)
930 931 932 933 934 935 936
                node.body[0:0] = self.tuple_to_stmts(origin_inner_tuple_node,
                                                     inner_tuple_name)
        elif self.is_for_iter(node) and isinstance(node.target,
                                                   (gast.List, gast.Tuple)):
            # Non-enumrate case:
            tuple_name = unique_name.generate(FOR_ITER_TUPLE_PREFIX)
            origin_tuple_node = node.target
937 938 939 940
            node.target = gast.Name(id=tuple_name,
                                    ctx=gast.Store(),
                                    annotation=None,
                                    type_comment=None)
941 942 943 944 945
            node.body[0:0] = self.tuple_to_stmts(origin_tuple_node, tuple_name)
        return node

    def tuple_to_stmts(self, node, tuple_name, idx=[]):
        if not isinstance(node, (gast.Tuple, gast.List)):
946
            value_node_str = tuple_name
947
            for i in idx:
948 949 950 951 952 953 954
                value_node_str = value_node_str + "[{}]".format(i)

            node_str = ast_to_source_code(node).strip()
            assign_node_str = "{} = {}".format(node_str, value_node_str)
            assign_node = gast.parse(assign_node_str).body[0]
            return [assign_node]

955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982
        # isinstance(node, (gast.Tuple, gast.List))
        ret = []
        for i, element in enumerate(node.elts):
            ret += self.tuple_to_stmts(node.elts[i], tuple_name, idx + [i])
        return ret

    def is_for_iter(self, for_node):
        assert isinstance(for_node,
                          gast.For), "Input node is not gast.For node."
        if isinstance(for_node.iter, (gast.Name, gast.Attribute)):
            return True
        elif isinstance(for_node.iter, gast.Call) and isinstance(
                for_node.iter.func,
                gast.Attribute) and for_node.iter.func.attr == 'numpy':
            return True
        elif isinstance(for_node.iter, gast.Subscript):
            return True
        else:
            return False

    def is_for_enumerate_iter(self, for_node):
        assert isinstance(for_node,
                          gast.For), "Input node is not gast.For node."
        return isinstance(for_node.iter, gast.Call) and isinstance(
            for_node.iter.func,
            gast.Name) and for_node.iter.func.id == "enumerate"


983
class ForNodeVisitor(object):
984
    """
985
    This class parses python for statement, get transformed 3 statement components of for node
986 987 988 989 990 991 992 993
    three key statements:
        1). init_stmts: list[node], prepare nodes of for loop, may not only one
        2). cond_stmt: node, condition node to judge whether continue loop
        3). body_stmts: list[node], updated loop body, sometimes we should change
            the original statement in body, not just append new statement

    In this process, the semantics of for does not change.

994
    Now only can parse 3 type statements (Here var is VarBase(Tensor) or python variable):
995 996 997
        1). for x in range(var[*]|var.numpy()[*])
        2). for x in var|var.numpy()
        3). for i, x enumerate(var|var.numpy())
998 999 1000 1001 1002
    """

    def __init__(self, for_node):
        assert isinstance(
            for_node, gast.For
1003
        ), "Input node for the initialization of ForNodeVisitor is not gast.For node."
1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016
        # 1. original for node
        self.node = for_node

        # 2. gast.For node main parts
        self.target = for_node.target
        # NOTE: type may be Node or list[Node]
        self.iter_args = for_node.iter if self.is_for_iter(
        ) else for_node.iter.args
        self.body = for_node.body

        # 3. key shared node or names
        # - x:
        #   - for x in range(***)
1017 1018
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
1019 1020 1021
        self.iter_var_name = self._get_iter_var_name()

        # - created index var to slice Variable: __for_loop_var_index_0
1022 1023
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
1024 1025
        self.iter_idx_name = unique_name.generate(FOR_ITER_INDEX_PREFIX)

1026
        # - created shape var to build loop condition: __for_loop_var_len_0
1027 1028 1029
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
        #   - for x in var
1030
        self.iter_var_len_name = unique_name.generate(FOR_ITER_VAR_LEN_PREFIX)
1031 1032 1033
        # - created zip to list var : __for_loop_iter_zip_0
        self.iter_zip_to_list_name = unique_name.generate(
            FOR_ITER_ZIP_TO_LIST_PREFIX)
1034

1035 1036 1037
        # - var.numpy()/var
        #   - for x in var|var.numpy()
        #   - for i, x enumerate(var|var.numpy())
1038 1039 1040
        self.iter_node = self._get_iter_node()

        # - enumeate i:
1041
        #   - for i, x enumerate(var|var.numpy())
1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055
        self.enum_idx_name = self._get_enum_idx_name()

        # - range/enumerate args length
        self.args_length = None

    def parse(self):
        self._args_check()
        if self.is_for_range_iter():
            return self._parse_for_range_stmts()
        elif self.is_for_iter():
            return self._parse_for_stmts()
        elif self.is_for_enumerate_iter():
            return self._parse_for_enumerate_stmts()
        else:
1056
            return None
1057 1058

    def is_for_range_iter(self):
1059 1060 1061
        return isinstance(self.node.iter, gast.Call) and isinstance(
            self.node.iter.func,
            gast.Name) and self.node.iter.func.id == "range"
1062 1063

    def is_for_iter(self):
1064 1065
        if isinstance(self.node.iter,
                      (gast.Name, gast.Attribute, gast.List, gast.Tuple)):
1066 1067 1068 1069 1070
            return True
        elif isinstance(self.node.iter, gast.Call) and isinstance(
                self.node.iter.func,
                gast.Attribute) and self.node.iter.func.attr == 'numpy':
            return True
1071 1072
        elif isinstance(self.node.iter, gast.Subscript):
            return True
1073 1074
        else:
            return False
1075 1076

    def is_for_enumerate_iter(self):
1077 1078 1079
        return isinstance(self.node.iter, gast.Call) and isinstance(
            self.node.iter.func,
            gast.Name) and self.node.iter.func.id == "enumerate"
1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105

    def _args_check(self):
        if self.is_for_range_iter():
            self.args_length = len(self.iter_args)
            assert self.args_length >= 1 and self.args_length <= 3, "range() function takes 1 to 3 arguments"
        elif self.is_for_enumerate_iter():
            self.args_length = len(self.iter_args)
            assert self.args_length >= 1 and self.args_length <= 2, "enumerate() function takes 1 to 2 arguments"
        else:
            self.args_length = None

    def _parse_for_range_stmts(self):
        init_stmts = []
        init_stmts.append(self._build_index_init_node())

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body
        body_stmts.append(self._build_index_increase_node(step_node))

        return init_stmts, cond_stmt, body_stmts

    def _parse_for_stmts(self):
        init_stmts = []
1106
        init_stmts.extend(self._build_iter_node())
1107
        init_stmts.append(self._build_index_init_node())
1108
        init_stmts.append(self._build_var_len_assign_node())
1109 1110 1111 1112 1113 1114

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body
1115 1116 1117 1118 1119

        # NOTE(liym27): Here add a gast.Assign, and the target of it is gast.Name.
        # In NameNodeReplaceTransformer, using gast.Name to replace gast.Name is safe.
        target_node, assign_node = self._build_assign_var_slice_node()
        body_stmts[0:0] = [assign_node]
1120 1121
        for body_node in body_stmts:
            NameNodeReplaceTransformer(body_node, self.iter_var_name,
1122
                                       target_node)
1123 1124 1125 1126 1127 1128
        body_stmts.append(self._build_index_increase_node(step_node))

        return init_stmts, cond_stmt, body_stmts

    def _parse_for_enumerate_stmts(self):
        init_stmts = []
1129
        init_stmts.extend(self._build_iter_node())
1130
        init_stmts.append(self._build_index_init_node())
1131
        init_stmts.append(self._build_var_len_assign_node())
1132 1133 1134 1135 1136 1137 1138
        init_stmts.append(self._build_enum_init_node())

        compare_node = self._build_compare_node()
        step_node = self._build_step_node()
        cond_stmt = self._build_cond_stmt(step_node, compare_node)

        body_stmts = self.body
1139 1140 1141

        target_node, assign_node = self._build_assign_var_slice_node()
        body_stmts[0:0] = [assign_node]
1142 1143
        for body_node in body_stmts:
            NameNodeReplaceTransformer(body_node, self.iter_var_name,
1144 1145
                                       target_node)

1146 1147 1148 1149 1150 1151 1152 1153
        body_stmts.append(self._build_index_increase_node(step_node))
        body_stmts.append(self._build_enum_increase_node())

        return init_stmts, cond_stmt, body_stmts

    def _build_index_init_node(self):
        if self.is_for_range_iter():
            if self.args_length == 1:
1154
                index_init_value_str = '0'
1155
            else:
1156 1157
                index_init_value_str = ast_to_source_code(
                    self.iter_args[0]).strip()
1158 1159

            index_init_var_name = self.iter_var_name
1160
        else:
1161 1162 1163 1164 1165 1166 1167 1168
            index_init_value_str = '0'
            index_init_var_name = self.iter_idx_name

        index_init_node_source_str = "{target} = {value}".format(
            target=index_init_var_name, value=index_init_value_str)

        index_init_node = gast.parse(index_init_node_source_str).body[0]

1169 1170
        return index_init_node

1171 1172 1173 1174 1175
    def _build_var_len_assign_node(self):
        # get the length of iterable variable
        if isinstance(self.iter_node, gast.Call) and isinstance(
                self.iter_node.func,
                gast.Attribute) and self.iter_node.func.attr == 'numpy':
1176 1177
            iter_var_name = ast_to_source_code(
                self.iter_node.func.value).strip()
1178
        else:
1179 1180
            iter_var_name = ast_to_source_code(self.iter_node).strip()

1181
        convert_len_node_source_str = '{} = _jst.convert_len({})'.format(
1182 1183 1184 1185 1186
            self.iter_var_len_name, iter_var_name)

        convert_len_node = gast.parse(convert_len_node_source_str).body[0]

        return convert_len_node
1187

1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207
    def _build_iter_node(self):
        """
        Process special cases for iter_node inclue:
          - Case 1 (for zip):
            
            - for i, val in enumerate(zip(x, y))  # original code:
            
            - __for_loop_iter_zip_0 = list(zip(x, y))
            - for i, val in enumerate(__for_loop_iter_zip_0)
        """
        new_nodes = []
        if isinstance(self.iter_node, gast.Call) and isinstance(
                self.iter_node.func, gast.Name):
            if self.iter_node.func.id == 'zip':
                iter_var_name = ast_to_source_code(self.iter_node).strip()
                zip_to_list_str = "{target} = list({value})".format(
                    target=self.iter_zip_to_list_name, value=iter_var_name)
                zip_to_list_node = gast.parse(zip_to_list_str).body[0]
                new_nodes.append(zip_to_list_node)

1208 1209 1210 1211
                self.iter_node = gast.Name(id=self.iter_zip_to_list_name,
                                           ctx=gast.Load(),
                                           annotation=None,
                                           type_comment=None)
1212 1213 1214

        return new_nodes

1215 1216
    def _build_enum_init_node(self):
        if self.is_for_enumerate_iter() and self.args_length != 1:
1217 1218 1219 1220 1221 1222 1223
            init_value_str = ast_to_source_code(self.iter_args[1]).strip()
        else:
            init_value_str = '0'

        enum_init_node_source_str = "{} = {}".format(self.enum_idx_name,
                                                     init_value_str)
        enum_init_node = gast.parse(enum_init_node_source_str).body[0]
1224 1225 1226 1227 1228 1229 1230
        return enum_init_node

    def _build_compare_node(self):
        if self.is_for_range_iter():
            compare_node = self.iter_args[
                0] if self.args_length == 1 else self.iter_args[1]
        else:
1231 1232 1233 1234
            compare_node = gast.Name(id=self.iter_var_len_name,
                                     ctx=gast.Load(),
                                     annotation=None,
                                     type_comment=None)
1235 1236 1237 1238 1239
        return compare_node

    def _build_step_node(self):
        if self.is_for_range_iter():
            step_node = self.iter_args[
1240 1241
                2] if self.args_length == 3 else gast.Constant(value=1,
                                                               kind=None)
1242 1243 1244 1245 1246
        else:
            step_node = gast.Constant(value=1, kind=None)
        return step_node

    def _build_cond_stmt(self, step_node, compare_node):
1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257
        if not isinstance(step_node, (gast.Constant, gast.UnaryOp)):
            raise NotImplementedError(
                "Dynamic-to-Static only supports the step value is a constant or negative constant in 'for-range' statements, "
                "such as '2', '-3'. But received: '{}'. Please fix code to be compatible with Dynamic-to-Static."
                .format(ast_to_source_code(step_node).strip()))

        if isinstance(step_node, gast.UnaryOp) or step_node.value < 0:
            # eg:
            # range(max, min, -2)
            # ->
            # i > min
1258 1259 1260 1261 1262 1263 1264 1265
            return gast.Compare(left=gast.Name(
                id=self.iter_var_name
                if self.is_for_range_iter() else self.iter_idx_name,
                ctx=gast.Load(),
                annotation=None,
                type_comment=None),
                                ops=[gast.Gt()],
                                comparators=[compare_node])
1266 1267 1268 1269 1270
        else:
            # eg:
            # range(min, max, 2)
            # ->
            # i < max
1271
            return gast.Compare(left=gast.Name(
1272 1273
                id=self.iter_var_name
                if self.is_for_range_iter() else self.iter_idx_name,
1274
                ctx=gast.Load(),
1275 1276
                annotation=None,
                type_comment=None),
1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288
                                ops=[gast.Lt()],
                                comparators=[compare_node])

    def _build_index_increase_node(self, step_node):
        return gast.AugAssign(target=gast.Name(
            id=self.iter_var_name
            if self.is_for_range_iter() else self.iter_idx_name,
            ctx=gast.Store(),
            annotation=None,
            type_comment=None),
                              op=gast.Add(),
                              value=step_node)
1289

1290
    def _build_assign_var_slice_node(self):
1291 1292 1293
        var_slice_str = "{}[{}]".format(
            ast_to_source_code(self.iter_node).strip(), self.iter_idx_name)
        var_slice_node = gast.parse(var_slice_str).body[0].value
1294 1295 1296 1297
        new_iter_var_name = unique_name.generate(FOR_ITER_VAR_NAME_PREFIX)
        target_node, assign_node = create_assign_node(new_iter_var_name,
                                                      var_slice_node)
        return target_node, assign_node
1298 1299

    def _build_enum_increase_node(self):
1300 1301 1302 1303 1304 1305
        return gast.AugAssign(target=gast.Name(id=self.enum_idx_name,
                                               ctx=gast.Store(),
                                               annotation=None,
                                               type_comment=None),
                              op=gast.Add(),
                              value=gast.Constant(value=1, kind=None))
1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326

    def _get_iter_var_name(self):
        if self.is_for_range_iter():
            return self.target.id
        elif self.is_for_iter():
            return self.target.id
        elif self.is_for_enumerate_iter():
            return self.target.elts[1].id
        return None

    def _get_iter_node(self):
        if self.is_for_iter():
            return self.iter_args
        elif self.is_for_enumerate_iter():
            return self.iter_args[0]
        return None

    def _get_enum_idx_name(self):
        if self.is_for_enumerate_iter():
            return self.target.elts[0].id
        return None
1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408


class SplitAssignTransformer(gast.NodeTransformer):
    """
    This class transforms sequence assignments and multi-target assignments to normal assignments.
    """

    def __init__(self, ast_node):
        assert isinstance(ast_node, gast.AST)
        self.ast_root = ast_node

    def transform(self):
        self.visit(self.ast_root)

    def visit_Assign(self, node):
        target_nodes = node.targets
        if len(target_nodes) == 1:
            node = self._parse_sequence_assign(node)
        else:
            node = self._parse_multi_target_assign(node)
        return node

    def _parse_sequence_assign(self, node):
        """
        a, b = c, d
        ->
        a = c
        b = d
        """
        assert isinstance(node, gast.Assign)

        target_nodes = node.targets
        value_node = node.value
        if not isinstance(target_nodes[0], (gast.List, gast.Tuple)):
            return node
        if not isinstance(value_node, (gast.List, gast.Tuple)):
            return node

        targets = node.targets[0].elts
        values = node.value.elts
        if len(targets) != len(values):
            return node

        new_nodes = []
        for target, value in zip(targets, values):
            assign_node = gast.Assign(targets=[target], value=value)
            new_nodes.append(assign_node)

        return new_nodes

    def _parse_multi_target_assign(self, node):
        """
         Example 1:
         a = b = c
         ->
         b = c
         a = b

         Example 2:
         a, b = c, d = x
         ->
         c,d = x
         a = c
         b = d
         """
        assert isinstance(node, gast.Assign)

        target_nodes = node.targets
        value_node = node.value
        new_nodes = []
        for target in reversed(target_nodes):
            assign_node = gast.Assign(targets=[target], value=value_node)
            # NOTE: Because assign_node can be sequence assign statement like `a,b = c,d`,
            # it's necessary to visit this new assign_node
            parsed_node = self.visit_Assign(assign_node)
            if not isinstance(parsed_node, list):
                parsed_node = [parsed_node]

            new_nodes.extend(parsed_node)
            value_node = target

        return new_nodes
1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424


# NOTE: inspect.unwrap() exits in PY3 but not in PY2.
def unwrap(func):
    """
    Returns the object wrapped by decorators.
    """

    def _is_wrapped(f):
        return hasattr(f, '__wrapped__')

    unwrapped_f = func
    while (_is_wrapped(unwrapped_f)):
        unwrapped_f = unwrapped_f.__wrapped__

    return unwrapped_f
1425 1426


C
Chen Weihang 已提交
1427
def input_specs_compatible(src_input_specs, desired_input_specs):
1428 1429 1430 1431
    """
    Returns True if the two input specs are compatible, otherwise False.

    args:
1432 1433 1434 1435
        src_input_spec (list or tuple[InputSpec et.al]): list/tuple of
            paddle.static.InputSpec or int/str et.al
        desired_input_specs (list or tuple[InputSpec et.al]): list/tuple of
            paddle.static.InputSpec or int/str et.al
1436 1437
    """
    len_specs = len(src_input_specs)
C
Chen Weihang 已提交
1438 1439
    if len_specs != len(desired_input_specs):
        # NOTE(chenweihang): if the input_spec of jit.save is a subset of
1440
        # input_spec of to_static, also compatible
C
Chen Weihang 已提交
1441 1442 1443 1444
        for spec in src_input_specs:
            if spec not in desired_input_specs:
                return False
    else:
1445 1446 1447 1448 1449 1450 1451 1452
        for (src_spec, desired_spec) in zip(src_input_specs,
                                            desired_input_specs):
            if isinstance(src_spec, paddle.static.InputSpec) or isinstance(
                    desired_spec, paddle.static.InputSpec):
                if not _compatible_tensor_spec(src_spec, desired_spec):
                    return False
            else:
                if not _compatible_non_tensor_spec(src_spec, desired_spec):
C
Chen Weihang 已提交
1453 1454
                    return False

1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481
    return True


def _compatible_tensor_spec(src_spec, desired_spec):
    """
    Check whether two tensor type spec is compatible.
    """
    for spec in [src_spec, desired_spec]:
        if not isinstance(spec, paddle.static.InputSpec):
            return False
    src_shape = src_spec.shape
    other_shape = desired_spec.shape
    len_shape = len(src_shape)
    if len_shape != len(other_shape):
        return False
    for j in range(len_shape):
        if src_shape[j] is None or src_shape[j] < 0:
            continue
        if other_shape[j] is None or other_shape[j] < 0:
            continue
        if src_shape[j] != other_shape[j]:
            return False

    src_dtype = convert_dtype(src_spec.dtype)
    other_dtype = convert_dtype(desired_spec.dtype)
    if src_dtype != other_dtype:
        return False
1482 1483

    return True
1484

1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505

def _compatible_non_tensor_spec(src_spec, desired_spec):
    """
    Check whether two non-tensor type spec is compatible.
    """

    def hash_value(spec):
        try:
            hash_val = make_hashable(spec)
        except:
            hash_val = None
        return hash_val

    src_hash_val = hash_value(src_spec)
    desired_hash_val = hash_value(desired_spec)

    if src_hash_val != desired_hash_val:
        return False
    else:
        return True

1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529

def slice_is_num(slice_node):
    # A slice_node.slice can be a:
    # (1) ast.Index, which is a simple number such as [1], [-2]
    # (2) ast.Slice, which is represented by bounds such as [2:-1]
    # (3) ast.Tuple, which includes the above two cases such as [2:-1, 1]
    # If slice node is case (1), return True, Otherwise, return False.
    #
    # NOTE: In (1) case, when gast>=0.4.0, gast.Index is not used, which is replaced
    # other gast node such as gast.Constant, gast.Name, gast.UnaryOp and so on.
    # Considering the compatibility of gast, here use ast note to check whether the
    # node is a num. For more details, please visit https://github.com/serge-sans-paille/gast

    assert isinstance(slice_node, gast.Subscript)
    slice_node_str = ast_to_source_code(slice_node).strip()
    ast_node = ast.parse(slice_node_str).body[0].value

    if isinstance(ast_node.slice, (ast.Tuple, ast.Slice)):
        return False

    if isinstance(ast_node.slice, ast.Index):
        return True

    return False