prim2code.py 29.9 KB
Newer Older
S
SunAhong1993 已提交
1
# -*- coding:UTF-8 -*-
S
SunAhong1993 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15
#   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.

S
SunAhong1993 已提交
16
NO_OUTPUT_COUNT = 0
S
SunAhong1993 已提交
17 18


S
SunAhong1993 已提交
19 20 21 22 23 24 25 26 27 28 29
def gen_codes(code_list, indent=0):
    indent_blank = "    " * indent
    codes = []
    for code_line in code_list:
        if code_line.strip() == "":
            codes.append('\n')
        else:
            codes.append(indent_blank + code_line + '\n')
    return codes


S
SunAhong1993 已提交
30
def get_value(layer, key, layer_id=None, different_attrs=None):
S
SunAhong1993 已提交
31 32 33 34 35 36
    """ 进行optimizer后可能把inputs的value直接用数值代替(ConstantFuser),
        会把input换成attr,所以需要此处的操作。
    """
    if key in layer.inputs:
        return layer.inputs[key]
    else:
S
SunAhong1993 已提交
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
        if different_attrs is None:
            return str(layer.attrs[key])
        else:
            key_name = "{}_{}".format(layer.outputs[0], key)
            if key_name in different_attrs:
                return key_name
            else:
                if layer_id is None:
                    return str(layer.attrs[key])
                key_name = "{}_{}".format("layer_id/{}".format(layer_id), key)
                if key_name in different_attrs:
                    new_key_name = "attr_{}".format(NO_OUTPUT_COUNT)
                    NO_OUTPUT_COUNT += 1
                    diff_index = different_attrs.index(key_name)
                    different_attrs[diff_index] = new_key_name
                    return new_key_name
                else:
                    return str(layer.attrs[key])


S
SunAhong1993 已提交
57 58 59 60 61 62
def prim_add(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
S
SunAhong1993 已提交
63
    line = "{} = {} + {}".format(layer.outputs[0],
S
SunAhong1993 已提交
64 65
                                 get_value(layer, "x", different_attrs),
                                 get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
66 67 68
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
69 70 71 72 73 74
def prim_add_(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
W
wjj19950828 已提交
75 76 77 78 79 80 81 82 83
    if layer.attrs["alpha"] == 1:
        line = "{} = {} + {}".format(layer.outputs[0],
                                     get_value(layer, "x", different_attrs),
                                     get_value(layer, "y", different_attrs))
    else:
        line = "{} = {} + {} * {}".format(
            layer.outputs[0],
            get_value(layer, "x", different_attrs), layer.attrs["alpha"],
            get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
84 85 86
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
87 88 89 90 91 92 93
def prim_and(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None,
             is_return_line=False):
S
SunAhong1993 已提交
94
    line = "{} = {} and {}".format(layer.outputs[0],
S
SunAhong1993 已提交
95 96
                                   get_value(layer, "x", different_attrs),
                                   get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
97 98
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
99 100 101
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
102 103 104 105 106 107
def prim_append(layer,
                indent=1,
                init_func=[],
                forward_func=[],
                layer_id=None,
                different_attrs=None):
S
SunAhong1993 已提交
108
    line = "{}.append({})".format(
S
SunAhong1993 已提交
109
        get_value(layer, "list", layer_id, different_attrs),
S
SunAhong1993 已提交
110
        get_value(layer, "element", layer_id, different_attrs))
S
SunAhong1993 已提交
111 112 113
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
114 115 116 117 118 119
def prim_assert(layer,
                indent=1,
                init_func=[],
                forward_func=[],
                layer_id=None,
                different_attrs=None):
S
SunAhong1993 已提交
120 121 122 123 124 125 126 127 128 129
    if layer.attrs["type"] == "eq":
        values = get_value(layer, "key")
        if "value" in layer.attrs:
            values = layer.attrs["value"]
        if isinstance(values, list):
            s = ""
            for v in values:
                s += "{} == {} or ".format(get_value(layer, "key"), v)
            if len(s) > 0:
                s = s[:-4]
S
SunAhong1993 已提交
130
            lc = locals()
S
SunAhong1993 已提交
131 132
            exec("assert_result = {}".format(s))
            assert_result = lc['assert_result']
S
SunAhong1993 已提交
133 134 135
            line = "assert {}, \'The {} must be {}!\'".format(
                s, get_value(layer, "key"), get_value(layer, "value"))
        else:
S
SunAhong1993 已提交
136 137 138
            s = "{} == {}".format(
                get_value(layer, "key"), get_value(layer, "value"))
            lc = locals()
S
SunAhong1993 已提交
139 140 141 142
            exec("assert_result = {}".format(s))
            assert_result = lc['assert_result']
            line = "assert {}, \'The {} must be {}!\'".format(
                s, get_value(layer, "key"), get_value(layer, "value"))
S
SunAhong1993 已提交
143 144
    else:
        raise Exception("Not implement yet!")
S
SunAhong1993 已提交
145 146
    if not assert_result:
        forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
147 148


S
SunAhong1993 已提交
149 150 151 152 153 154
def prim_check_dim(layer,
                   indent=1,
                   init_func=[],
                   forward_func=[],
                   layer_id=None,
                   different_attrs=None):
S
SunAhong1993 已提交
155
    lines = []
S
SunAhong1993 已提交
156 157
    dim = get_value(layer, "dim", different_attrs)
    lines.append("if {} < 0:".format(dim))
S
SunAhong1993 已提交
158
    lines.append("    {} = {} + {}".format(layer.outputs[
S
SunAhong1993 已提交
159
        0], dim, get_value(layer, "len", different_attrs)))
S
SunAhong1993 已提交
160
    lines.append("else:")
S
SunAhong1993 已提交
161
    lines.append("    {} = {}".format(layer.outputs[0], dim))
S
SunAhong1993 已提交
162 163 164
    forward_func.extend(gen_codes(lines, indent=indent))


S
SunAhong1993 已提交
165 166 167 168 169 170
def prim_constant(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
S
SunAhong1993 已提交
171 172 173 174
    line = "{} = {}".format(layer.outputs[0], layer.attrs["value"])
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
175 176 177 178 179 180 181
def prim_contain(layer,
                 indent=1,
                 init_func=[],
                 forward_func=[],
                 layer_id=None,
                 different_attrs=None,
                 is_return_line=False):
S
SunAhong1993 已提交
182
    line = "{} = {} in {}".format(layer.outputs[0],
S
SunAhong1993 已提交
183 184
                                  get_value(layer, "element", different_attrs),
                                  get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
185 186
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
187 188 189
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
190 191 192 193 194 195
def prim_dict(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
S
SunAhong1993 已提交
196 197
    line = "{} = dict()".format(layer.outputs[0])
    forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
198 199 200 201 202 203 204 205


def prim_dict_construct(layer,
                        indent=1,
                        init_func=[],
                        forward_func=[],
                        layer_id=None,
                        different_attrs=None):
S
SunAhong1993 已提交
206
    lines = list()
S
SunAhong1993 已提交
207
    line = "{} = dict()".format(layer.outputs[0])
S
SunAhong1993 已提交
208 209
    lines.append(line)
    for i in range(len(layer.inputs)):
S
SunAhong1993 已提交
210 211 212 213
        line = "{}[{}] = {}".format(
            layer.outputs[0],
            get_value(layer, "key{}".format(i), different_attrs),
            get_value(layer, "value{}".format(i), different_attrs))
S
SunAhong1993 已提交
214 215
        lines.append(line)
    forward_func.extend(gen_codes(lines, indent=indent))
S
SunAhong1993 已提交
216 217 218 219 220 221 222 223 224 225


def prim_dict2values(layer,
                     indent=1,
                     init_func=[],
                     forward_func=[],
                     layer_id=None,
                     different_attrs=None):
    line = "{} = list({}.values())".format(
        layer.outputs[0], get_value(layer, "x", different_attrs))
S
SunAhong1993 已提交
226 227 228
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
229 230 231 232 233 234
def prim_div(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
S
SunAhong1993 已提交
235
    line = "{} = {} / {}".format(layer.outputs[0],
S
SunAhong1993 已提交
236
                                 get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
237
                                 get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
238 239 240
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
241 242 243 244 245 246 247
def prim_eq(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
248
    line = "{} = {} == {}".format(layer.outputs[0],
S
SunAhong1993 已提交
249
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
250
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
251 252
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
253 254 255
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
256 257 258 259 260 261 262 263
def prim_equal(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
    line = "{} = {}".format(layer.outputs[0],
                            get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
264 265 266
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
267 268 269 270 271 272 273 274
def prim_exception(layer,
                   indent=1,
                   init_func=[],
                   forward_func=[],
                   layer_id=None,
                   different_attrs=None):
    line = "raise Exception({})".format(
        get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
275 276 277
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
278 279 280 281 282 283 284 285
def prim_float(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
    line = "{} = float({})".format(layer.outputs[0],
                                   get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
286 287 288
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
289 290 291 292 293 294
def prim_floor(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
S
SunAhong1993 已提交
295
    line = "{} = math.floor({})".format(layer.outputs[0],
S
SunAhong1993 已提交
296
                                        get_value(layer, "x", different_attrs))
S
SunAhong1993 已提交
297 298 299
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
300 301 302 303 304 305
def prim_floordiv(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
S
SunAhong1993 已提交
306
    line = "{} = {} // {}".format(layer.outputs[0],
S
SunAhong1993 已提交
307
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
308
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
309 310 311
    forward_func.extend(gen_codes([line], indent=indent))


312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332
def prim_format(layer,
                indent=1,
                init_func=[],
                forward_func=[],
                layer_id=None,
                different_attrs=None):
    line = ""
    if len(layer.inputs) == 3:
        line = "{} = {}.format({}, {})".format(
            layer.outputs[0],
            get_value(layer, "input0", different_attrs),
            get_value(layer, "input1", different_attrs),
            get_value(layer, "input2", different_attrs))
    elif len(layer.inputs) == 2:
        line = "{} = {}.format({})".format(
            layer.outputs[0],
            get_value(layer, "input0", different_attrs),
            get_value(layer, "input1", different_attrs))
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
333 334 335 336 337 338
def prim_getitem(layer,
                 indent=1,
                 init_func=[],
                 forward_func=[],
                 layer_id=None,
                 different_attrs=None):
S
SunAhong1993 已提交
339
    line = "{} = {}[{}]".format(layer.outputs[0],
S
SunAhong1993 已提交
340 341
                                get_value(layer, "list", different_attrs),
                                get_value(layer, "index", different_attrs))
S
SunAhong1993 已提交
342 343 344
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
345 346 347 348 349 350 351
def prim_gt(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
352
    line = "{} = {} > {}".format(layer.outputs[0],
S
SunAhong1993 已提交
353
                                 get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
354
                                 get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
355 356
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
357 358 359
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
360 361 362 363 364 365
def prim_if(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None):
S
SunAhong1993 已提交
366 367 368 369
    try:
        exec_s = None
        for line in forward_func:
            s = line.replace("    ", "")
S
SunAhong1993 已提交
370 371
            if s.startswith("{} = ".format(
                    get_value(layer, "input", different_attrs))):
S
SunAhong1993 已提交
372
                exec_s = s.split(" = ")[1]
S
SunAhong1993 已提交
373
        lc = locals()
S
SunAhong1993 已提交
374 375 376
        if exec_s is not None:
            exec("if_result = {}".format(exec_s))
        else:
S
SunAhong1993 已提交
377 378
            exec("if_result = {}".format(
                get_value(layer, "input", different_attrs)))
S
SunAhong1993 已提交
379 380 381 382 383 384
        if_result = lc['if_result']
        if if_result:
            block = layer.blocks[0]
        else:
            block = layer.blocks[1]
        if len(block.layers) > 0:
S
SunAhong1993 已提交
385
            b_init_lines, b_forward_lines = block.gen_code(indent=indent)
S
SunAhong1993 已提交
386 387 388 389 390 391 392 393 394 395
            init_func.extend(b_init_lines)
            forward_func.extend(b_forward_lines)
    except:
        line = "if {} :".format(get_value(layer, "input", different_attrs))
        forward_func.extend(gen_codes([line], indent=indent))
        block = layer.blocks[0]
        if len(block.layers) == 0:
            line = "pass"
            forward_func.extend(gen_codes([line], indent=indent + 1))
        else:
S
SunAhong1993 已提交
396
            b_init_lines, b_forward_lines = block.gen_code(indent=indent + 1)
S
SunAhong1993 已提交
397 398 399 400
            init_func.extend(b_init_lines)
            forward_func.extend(b_forward_lines)
        block = layer.blocks[1]
        if len(block.layers) > 0:
S
SunAhong1993 已提交
401
            b_init_lines, b_forward_lines = block.gen_code(indent=indent + 1)
S
SunAhong1993 已提交
402 403 404 405 406
            if len(b_forward_lines) != 0:
                line = "else:"
                forward_func.extend(gen_codes([line], indent=indent))
            init_func.extend(b_init_lines)
            forward_func.extend(b_forward_lines)
S
SunAhong1993 已提交
407 408


S
SunAhong1993 已提交
409 410 411 412 413 414 415 416
def prim_int(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
    line = "{} = int({})".format(layer.outputs[0],
                                 get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
417 418 419
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
420 421 422 423 424 425 426
def prim_is(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
427
    line = "{} = {} is {}".format(layer.outputs[0],
S
SunAhong1993 已提交
428
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
429
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
430 431
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
432 433 434
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
435 436 437 438 439 440 441 442 443 444
def prim_isinstance(layer,
                    indent=1,
                    init_func=[],
                    forward_func=[],
                    layer_id=None,
                    different_attrs=None,
                    is_return_line=False):
    line = "{} = isinstance({}, {})".format(
        layer.outputs[0],
        get_value(layer, "input", different_attrs), layer.attrs["cls"])
S
SunAhong1993 已提交
445 446
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
447 448 449
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
450 451 452 453 454 455 456
def prim_isnot(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None,
               is_return_line=False):
S
SunAhong1993 已提交
457
    line = "{} = {} is not {}".format(layer.outputs[0],
S
SunAhong1993 已提交
458 459
                                      get_value(layer, "x", different_attrs),
                                      get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
460 461
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
462 463 464
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
465 466 467 468 469 470 471
def prim_le(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
472
    line = "{} = {} <= {}".format(layer.outputs[0],
S
SunAhong1993 已提交
473
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
474
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
475 476
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
477 478 479
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
480 481 482 483 484 485 486 487
def prim_len(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
    line = "{} = len({})".format(layer.outputs[0],
                                 get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
488 489 490
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
491 492 493 494 495 496
def prim_len2list(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
S
SunAhong1993 已提交
497 498
    lines = []
    lines.append("{} = []".format(layer.outputs[0]))
S
SunAhong1993 已提交
499 500
    lines.append("for i in range({}):".format(
        get_value(layer, "len", different_attrs)))
S
SunAhong1993 已提交
501 502 503 504
    lines.append("    {}.append(i)".format(layer.outputs[0]))
    forward_func.extend(gen_codes(lines, indent=indent))


S
SunAhong1993 已提交
505 506 507 508 509 510 511
def prim_lt(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
512
    line = "{} = {} < {}".format(layer.outputs[0],
S
SunAhong1993 已提交
513
                                 get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
514
                                 get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
515 516
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
517 518 519
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
520 521 522 523 524 525
def prim_list(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
S
SunAhong1993 已提交
526 527 528
    input_len = len(layer.inputs) + len(layer.attrs)
    inputs_list = list()
    for i in range(input_len):
S
SunAhong1993 已提交
529 530
        inputs_list.append(
            get_value(layer, "input{}".format(i), different_attrs))
S
SunAhong1993 已提交
531 532 533 534 535
    inputs_str = ', '.join(inputs_list)
    line = "{} = [{}]".format(layer.outputs[0], inputs_str)
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
536 537 538 539 540 541 542 543
def prim_list_unpack(layer,
                     indent=1,
                     init_func=[],
                     forward_func=[],
                     layer_id=None,
                     different_attrs=None):
    line = "{} = {}".format(", ".join(layer.outputs),
                            get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
544 545 546
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
547 548 549 550 551 552
def prim_loop(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
S
SunAhong1993 已提交
553
    loop_range = get_value(layer, "input", different_attrs)
S
SunAhong1993 已提交
554 555 556
    line = "for {} in range({}):".format(layer.outputs[1], loop_range)
    forward_func.extend(gen_codes([line], indent=indent))
    block = layer.blocks[0]
S
SunAhong1993 已提交
557
    b_init_lines, b_forward_lines = block.gen_code(indent=indent + 1)
S
SunAhong1993 已提交
558 559 560 561
    init_func.extend(b_init_lines)
    forward_func.extend(b_forward_lines)


S
SunAhong1993 已提交
562 563 564 565 566 567 568 569
def prim_min(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
    line = "{} = min({})".format(layer.outputs[0],
                                 get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
570 571 572
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
573 574 575 576 577 578
def prim_mul(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
S
SunAhong1993 已提交
579
    line = "{} = {} * {}".format(layer.outputs[0],
S
SunAhong1993 已提交
580
                                 get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
581
                                 get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
582 583 584
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
585 586 587 588 589 590 591
def prim_ne(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
592
    line = "{} = {} != {}".format(layer.outputs[0],
S
SunAhong1993 已提交
593
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
594
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
595 596
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
597 598 599
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
600 601 602 603 604 605 606 607
def prim_neg(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
    line = "{} = -{}".format(layer.outputs[0],
                             get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
608 609 610
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
611 612 613 614 615 616 617 618 619
def prim_not(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None,
             is_return_line=False):
    line = "{} = not {}".format(layer.outputs[0],
                                get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
620 621
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
622 623 624
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
625 626 627 628 629 630 631
def prim_or(layer,
            indent=1,
            init_func=[],
            forward_func=[],
            layer_id=None,
            different_attrs=None,
            is_return_line=False):
S
SunAhong1993 已提交
632
    line = "{} = {} or {}".format(layer.outputs[0],
S
SunAhong1993 已提交
633
                                  get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
634
                                  get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
635 636
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
637
    forward_func.extend(gen_codes([line], indent=indent))
638 639


640 641 642 643 644 645 646 647
def prim_remainder(layer,
                   indent=1,
                   init_func=[],
                   forward_func=[],
                   layer_id=None,
                   different_attrs=None,
                   is_return_line=False):
    line = "{} = {} % {}".format(layer.outputs[0],
648 649
                                 get_value(layer, "x", different_attrs),
                                 get_value(layer, "y", different_attrs))
650 651 652
    if is_return_line:
        return line.split(" = ")[1]
    forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
653 654


S
SunAhong1993 已提交
655 656 657 658 659 660
def prim_replaceitem(layer,
                     indent=1,
                     init_func=[],
                     forward_func=[],
                     layer_id=None,
                     different_attrs=None):
S
SunAhong1993 已提交
661
    line = "{}[{}] = {}".format(
S
SunAhong1993 已提交
662
        get_value(layer, "list", layer_id, different_attrs),
S
SunAhong1993 已提交
663
        get_value(layer, "index", layer_id, different_attrs),
S
SunAhong1993 已提交
664
        get_value(layer, "item", layer_id, different_attrs))
S
SunAhong1993 已提交
665 666 667
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
668 669 670 671 672 673 674 675
def prim_requires_grad(layer,
                       indent=1,
                       init_func=[],
                       forward_func=[],
                       layer_id=None,
                       different_attrs=None):
    line = "{} = not {}.stop_gradient".format(
        layer.outputs[0], get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
676 677 678
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
679 680 681 682 683 684 685 686 687 688 689
def prim_rsub(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
    line = "{} = {} - {} * {}".format(
        layer.outputs[0],
        get_value(layer, "y", different_attrs),
        get_value(layer, "x", different_attrs),
        get_value(layer, "alpha", different_attrs))
S
SunAhong1993 已提交
690 691 692
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
693 694 695 696 697 698 699 700
def prim_select(layer,
                indent=1,
                init_func=[],
                forward_func=[],
                layer_id=None,
                different_attrs=None):
    line = "{} = {}[".format(layer.outputs[0],
                             get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
701 702
    for dim in range(layer.attrs["dim"]):
        line += ":, "
S
SunAhong1993 已提交
703
    line += (get_value(layer, "index", different_attrs) + "]")
S
SunAhong1993 已提交
704 705 706
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
707 708 709 710 711 712 713 714
def prim_set_attr(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
    line = "{} = {}".format(layer.outputs[0],
                            get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
715 716 717
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
718 719 720 721 722 723
def prim_set_item(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
S
SunAhong1993 已提交
724
    line = "{}[{}] = {}".format(
S
SunAhong1993 已提交
725
        get_value(layer, "dict", different_attrs),
S
SunAhong1993 已提交
726 727
        get_value(layer, "key", different_attrs),
        get_value(layer, "value", different_attrs))
S
SunAhong1993 已提交
728 729 730
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
731 732 733 734 735 736 737 738
def prim_shape(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
    line = "{} = {}.shape".format(layer.outputs[0],
                                  get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
739 740 741
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
742 743 744 745 746 747 748 749 750 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
def prim_shape_dim(layer,
                   indent=1,
                   init_func=[],
                   forward_func=[],
                   layer_id=None,
                   different_attrs=None):
    line = "{} = {}.shape[{}]".format(
        layer.outputs[0],
        get_value(layer, "input", different_attrs),
        get_value(layer, "dim", different_attrs))
    forward_func.extend(gen_codes([line], indent=indent))


def prim_slice(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
    line = "{} = {}[{}: {}: {}]".format(
        layer.outputs[0],
        get_value(layer, "input", different_attrs),
        get_value(layer, "start", different_attrs),
        get_value(layer, "end", different_attrs),
        get_value(layer, "step", different_attrs))
    forward_func.extend(gen_codes([line], indent=indent))


def prim_startswith(layer,
                    indent=1,
                    init_func=[],
                    forward_func=[],
                    layer_id=None,
                    different_attrs=None,
                    is_return_line=False):
    line = "{} = {}.startswith({})".format(
        layer.outputs[0],
        get_value(layer, "input", different_attrs),
        get_value(layer, "start_str", different_attrs))
S
fix3  
SunAhong1993 已提交
781 782
    if is_return_line:
        return line.split(" = ")[1]
S
SunAhong1993 已提交
783
    forward_func.extend(gen_codes([line], indent=indent))
S
SunAhong1993 已提交
784 785


S
SunAhong1993 已提交
786 787 788 789 790 791 792 793
def prim_str(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
    line = "{} = str({})".format(layer.outputs[0],
                                 get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
794 795 796
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
797 798 799 800 801 802
def prim_sub(layer,
             indent=1,
             init_func=[],
             forward_func=[],
             layer_id=None,
             different_attrs=None):
S
fix  
SunAhong1993 已提交
803
    if int(float(get_value(layer, "alpha", different_attrs))) == 1:
S
SunAhong1993 已提交
804
        line = "{} = {} - {}".format(layer.outputs[0],
S
SunAhong1993 已提交
805
                                     get_value(layer, "x", different_attrs),
S
SunAhong1993 已提交
806 807
                                     get_value(layer, "y", different_attrs))
    else:
S
SunAhong1993 已提交
808 809 810 811 812
        line = "{} = {} - {} * {}".format(
            layer.outputs[0],
            get_value(layer, "x", different_attrs),
            get_value(layer, "alpha", different_attrs),
            get_value(layer, "y", different_attrs))
S
SunAhong1993 已提交
813 814 815
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
816 817 818 819 820 821
def prim_tuple(layer,
               indent=1,
               init_func=[],
               forward_func=[],
               layer_id=None,
               different_attrs=None):
S
SunAhong1993 已提交
822 823 824
    input_len = len(layer.inputs) + len(layer.attrs)
    inputs_list = list()
    for i in range(input_len):
S
SunAhong1993 已提交
825 826
        inputs_list.append(
            get_value(layer, "input{}".format(i), different_attrs))
S
SunAhong1993 已提交
827 828 829 830 831
    inputs_str = ', '.join(inputs_list)
    line = "{} = ({})".format(layer.outputs[0], inputs_str)
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
832 833 834 835 836 837
def prim_tuple_unpack(layer,
                      indent=1,
                      init_func=[],
                      forward_func=[],
                      layer_id=None,
                      different_attrs=None):
S
SunAhong1993 已提交
838
    outputs_str = ', '.join(layer.outputs)
S
SunAhong1993 已提交
839 840
    line = "{} = {}".format(outputs_str,
                            get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
841 842 843
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
844 845 846 847 848 849 850 851
def prim_type(layer,
              indent=1,
              init_func=[],
              forward_func=[],
              layer_id=None,
              different_attrs=None):
    line = "{} = {}.dtype".format(layer.outputs[0],
                                  get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
852 853 854
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
855 856 857 858 859 860 861 862
def prim_var2list(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
    line = "{} = {}.numpy().tolist()".format(
        layer.outputs[0], get_value(layer, "input", different_attrs))
S
SunAhong1993 已提交
863 864 865
    forward_func.extend(gen_codes([line], indent=indent))


S
SunAhong1993 已提交
866 867 868 869 870 871
def prim_warnings(layer,
                  indent=1,
                  init_func=[],
                  forward_func=[],
                  layer_id=None,
                  different_attrs=None):
S
SunAhong1993 已提交
872 873
    lines = ["import warnings"]
    line = "warnings.warn({}, stacklevel={})".format(
S
SunAhong1993 已提交
874
        get_value(layer, "input", different_attrs), layer.attrs["stacklevel"])
S
SunAhong1993 已提交
875 876
    lines.append(line)
    forward_func.extend(gen_codes(lines, indent=indent))