control_flow.py 46.6 KB
Newer Older
Y
Yu Yang 已提交
1 2 3 4
from ..layer_helper import LayerHelper, unique_name
from ..framework import Program, Variable, Operator
from .. import core
from tensor import assign, fill_constant
D
dzhwinter 已提交
5
import contextlib
Y
Yang Yu 已提交
6
from ..registry import autodoc
D
dzhwinter 已提交
7

Q
QI JUN 已提交
8
__all__ = [
Y
Yang Yang 已提交
9 10 11 12 13
    'split_lod_tensor', 'merge_lod_tensor', 'BlockGuard',
    'BlockGuardWithCompletion', 'StaticRNNMemoryLink', 'WhileGuard', 'While',
    'lod_rank_table', 'max_sequence_len', 'topk', 'lod_tensor_to_array',
    'array_to_lod_tensor', 'increment', 'array_write', 'create_array',
    'less_than', 'array_read', 'shrink_memory', 'array_length', 'IfElse',
Y
Yang Yang 已提交
14 15
    'DynamicRNN', 'ConditionalBlock', 'StaticRNN', 'reorder_lod_tensor_by_rank',
    'ParallelDo'
D
dzhwinter 已提交
16 17
]

Y
Yu Yang 已提交
18

19
def split_lod_tensor(input, mask, level=0):
20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
    """
    **split_lod_tensor**

    This function takes in an input that contains the complete lod information,
    and takes in a mask which is used to mask certain parts of the input.
    The output is the true branch and the false branch with the mask applied to
    the input at a certain level in the tensor.

    Args:
        input(tuple|list|None): The input tensor that contains complete
                                lod information needed to construct the output.
        mask(list): A bool column vector which masks the input.
        level(int): The specific lod level to rank.

    Returns:
        Variable: The true branch of tensor as per the mask applied to input.
        Variable: The false branch of tensor as per the mask applied to input.

    Examples:
        .. code-block:: python

          x = layers.data(name='x', shape=[1])
          x.persistable = True

          y = layers.data(name='y', shape=[1])
          y.persistable = True

          out_true, out_false = layers.split_lod_tensor(
                input=x, mask=y, level=level)
    """
50
    helper = LayerHelper('split_lod_tensor', **locals())
F
fengjiayi 已提交
51 52
    out_true = helper.create_tmp_variable(dtype=input.dtype)
    out_false = helper.create_tmp_variable(dtype=input.dtype)
53 54 55 56 57 58 59 60 61 62 63 64
    helper.append_op(
        type='split_lod_tensor',
        inputs={
            'X': input,
            'Mask': mask,
        },
        outputs={'OutTrue': out_true,
                 'OutFalse': out_false},
        attrs={'level': level})
    return out_true, out_false


65
def merge_lod_tensor(in_true, in_false, x, mask, level=0):
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
    """
    **merge_lod_tensor**

    This function takes in an input :math:`x`, the True branch, the False
    branch and a binary :math:`mask`. Using this information, this function
    merges the True and False branches of the tensor into a single Output
    at a certain lod level indiacted by :math:`level`.

    Args:
        in_true(tuple|list|None): The True branch to be merged.
        in_false(tuple|list|None): The False branch to be merged.
        x(tuple|list|None): The input tensor that contains complete
                            lod information needed to construct the output.
        mask(list): A bool column vector which masks the input.
        level(int): The specific lod level to rank.

    Returns:
        Variable: The merged output tensor.

    Examples:
        .. code-block:: python

          x = layers.data(
                      name='x', shape=[1], dtype='float32', stop_gradient=False)
          y = layers.data(
                name='y', shape=[1], dtype='bool', stop_gradient=False)

          level = 0

          out_true, out_false = layers.split_lod_tensor(
                input=x, mask=y, level=level)
          out = layers.merge_lod_tensor(
                in_true=out_true, in_false=out_false, mask=y, x=x, level=level)
    """
100
    helper = LayerHelper('merge_lod_tensor', **locals())
F
fengjiayi 已提交
101
    out = helper.create_tmp_variable(dtype=in_true.dtype)
102 103 104 105 106 107 108 109 110 111 112
    helper.append_op(
        type='merge_lod_tensor',
        inputs={'X': x,
                'Mask': mask,
                'InTrue': in_true,
                'InFalse': in_false},
        outputs={'Out': out},
        attrs={'level': level})
    return out


Y
Yu Yang 已提交
113 114
class BlockGuard(object):
    """
115 116 117 118
    BlockGuard class.

    BlockGuard class is used to create a sub-block in a program by
    using the Python `with` keyword.
Y
Yu Yang 已提交
119 120
    """

121 122
    def __init__(self, main_program):
        if not isinstance(main_program, Program):
Y
Yu Yang 已提交
123
            raise TypeError("BlockGuard takes a program")
124
        self.main_program = main_program
Y
Yu Yang 已提交
125 126

    def __enter__(self):
127
        self.main_program.create_block()
Y
Yu Yang 已提交
128 129

    def __exit__(self, exc_type, exc_val, exc_tb):
130
        self.main_program.rollback()
Y
Yu Yang 已提交
131 132 133 134 135
        if exc_type is not None:
            return False  # re-raise exception
        return True


Y
Yang Yang 已提交
136
class ParallelDo(object):
137
    """
Y
Yang Yang 已提交
138
    ParallelDo class.
139

Y
Yang Yang 已提交
140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
    ParallelDo class is used to create a ParallelDo.
    """

    def __init__(self, places, name=None):
        self.helper = LayerHelper("parallel_do", name=name)
        self.inputs = []
        self.places = places
        self.outputs = []
        self.status = StaticRNN.BEFORE_RNN_BLOCK

    def do(self):
        return BlockGuardWithCompletion(self)

    def parent_block(self):
        prog = self.helper.main_program
        parent_idx = prog.current_block().parent_idx
        assert parent_idx >= 0
        parent_block = prog.block(parent_idx)
        return parent_block

    def __call__(self, *args, **kwargs):
        if self.status != StaticRNN.AFTER_RNN_BLOCK:
            raise ValueError("RNN output can only be retrieved after rnn block")
        if len(self.outputs) == 0:
            raise ValueError("RNN has no output")
        elif len(self.outputs) == 1:
            return self.outputs[0]
        else:
            return self.outputs

    def read_input(self, var):
        self.inputs.append(var)
Y
Yang Yang 已提交
172
        return var
Y
Yang Yang 已提交
173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208

    def write_output(self, var):
        self.outputs.append(var)

    def get_parameters(self):
        main_program = self.helper.main_program
        current_block = main_program.current_block()
        parent_block = self.parent_block()

        local_inputs = set()

        for op in current_block.ops:
            for oname in op.output_names:
                for out_var_name in op.output(oname):
                    local_inputs.add(out_var_name)

        for var in self.inputs:
            local_inputs.add(var.name)

        params = list()
        for op in current_block.ops:
            for iname in op.input_names:
                for in_var_name in op.input(iname):
                    if in_var_name not in local_inputs:
                        params.append(in_var_name)

        return [parent_block.var(name) for name in params]

    def complete_op(self):
        main_program = self.helper.main_program
        current_block = main_program.current_block()
        parent_block = self.parent_block()

        step_scope = parent_block.create_var(
            type=core.VarDesc.VarType.STEP_SCOPES)

Y
Yang Yang 已提交
209 210 211 212 213 214 215 216 217 218
        self.outputs = [
            parent_block.create_var(
                name=o.name,
                shape=o.shape,
                dtype=o.dtype,
                lod_level=o.lod_level,
                persistable=o.persistable,
                stop_gradient=o.stop_gradient) for o in self.outputs
        ]

Y
Yang Yang 已提交
219
        inputs = [parent_block.var(i.name) for i in self.inputs]
Y
Yang Yang 已提交
220
        outputs = [parent_block.var(o.name) for o in self.outputs]
Y
Yang Yang 已提交
221 222 223 224 225 226 227 228

        parent_block.append_op(
            type='parallel_do',
            inputs={
                'inputs': inputs,
                'parameters': self.get_parameters(),
                'places': self.places
            },
Y
Yang Yang 已提交
229
            outputs={'outputs': outputs,
Y
Yang Yang 已提交
230
                     'parallel_scopes': [step_scope]},
Y
Yang Yang 已提交
231 232 233 234 235 236 237 238
            attrs={'sub_block': current_block})


class BlockGuardWithCompletion(BlockGuard):
    """
    BlockGuardWithCompletion class.

    BlockGuardWithCompletion class is used to create an op with a block in a program.
239 240
    """

Y
Yu Yang 已提交
241
    def __init__(self, rnn):
Y
Yang Yang 已提交
242 243 244 245
        if not (isinstance(rnn, StaticRNN) or isinstance(rnn, ParallelDo)):
            raise TypeError(
                "BlockGuardWithCompletion takes a StaticRNN or ParallelDo")
        super(BlockGuardWithCompletion, self).__init__(rnn.helper.main_program)
Y
Yu Yang 已提交
246 247 248 249
        self.rnn = rnn

    def __enter__(self):
        self.rnn.status = StaticRNN.IN_RNN_BLOCK
Y
Yang Yang 已提交
250
        return super(BlockGuardWithCompletion, self).__enter__()
Y
Yu Yang 已提交
251 252

    def __exit__(self, exc_type, exc_val, exc_tb):
Y
Yu Yang 已提交
253 254
        if exc_type is not None:
            return False
Y
Yu Yang 已提交
255
        self.rnn.status = StaticRNN.AFTER_RNN_BLOCK
Y
Yang Yang 已提交
256 257 258
        self.rnn.complete_op()
        return super(BlockGuardWithCompletion, self).__exit__(exc_type, exc_val,
                                                              exc_tb)
Y
Yu Yang 已提交
259 260 261 262


class StaticRNNMemoryLink(object):
    """
263 264 265 266 267 268 269 270 271 272 273 274
    StaticRNNMemoryLink class.

    Args:
        init: the initial variable for Memory
        init: Variable
        pre_mem: the memory variable in previous time step
        pre_mem: Variable
        mem: the memory variable in current time step
        mem: Variable

    StaticRNNMemoryLink class is used to create a link between two
    memory cells of a StaticRNN.
Y
Yu Yang 已提交
275 276 277 278 279 280 281 282 283
    """

    def __init__(self, init, pre_mem, mem=None):
        self.init = init
        self.pre_mem = pre_mem
        self.mem = mem


class StaticRNN(object):
284 285 286 287 288 289
    """
    StaticRNN class.

    StaticRNN class is used to create a StaticRNN. The RNN will have its
    own parameters like inputs, outputs, memories, status and length.
    """
Y
Yu Yang 已提交
290 291 292 293
    BEFORE_RNN_BLOCK = 0
    IN_RNN_BLOCK = 1
    AFTER_RNN_BLOCK = 2

294 295
    def __init__(self, name=None):
        self.helper = LayerHelper("static_rnn", name=name)
Y
Yu Yang 已提交
296 297 298 299 300 301 302 303
        self.memories = {}  # memory map, from pre_mem.name --> MemoryLink
        self.inputs = []  # input variable list in current block
        self.outputs = []  # output variable list in parent block
        self.status = StaticRNN.BEFORE_RNN_BLOCK  # status flag.
        # sequence length, since it is a static RNN, sequence length are fixed.
        self.seq_len = None

    def step(self):
Y
Yang Yang 已提交
304
        return BlockGuardWithCompletion(self)
Y
Yu Yang 已提交
305 306 307 308 309

    def _assert_in_rnn_block_(self, method):
        if self.status != StaticRNN.IN_RNN_BLOCK:
            raise ValueError("You must invoke {0} in rnn block".format(method))

310 311 312 313 314 315 316
    def memory(self,
               init=None,
               shape=None,
               batch_ref=None,
               init_value=0.0,
               init_batch_dim_idx=0,
               ref_batch_dim_idx=1):
317 318 319 320 321 322 323 324 325
        """
        Args:
            init: boot memory, if not set, a shape, batch_ref must be provided
            shape: shape of the boot memory
            batch_ref: batch size reference variable
            init_value: the init value of boot memory
            init_batch_dim_idx: the index of batch size in init's dimension
            ref_batch_dim_idx: the index of batch size in batch_ref's dimension
        """
Y
Yu Yang 已提交
326 327
        self._assert_in_rnn_block_('memory')
        if init is None:
328
            if shape is None or batch_ref is None:
Y
Yu Yang 已提交
329
                raise ValueError(
330
                    "if init is None, memory at least need shape and batch_ref")
Y
Yu Yang 已提交
331 332 333
            parent_block = self.parent_block()
            var_name = unique_name("@".join([self.helper.name, "memory_boot"]))
            boot_var = parent_block.create_var(
334 335
                name=var_name,
                shape=shape,
F
fengjiayi 已提交
336
                dtype=batch_ref.dtype,
337
                persistable=False)
Y
Yu Yang 已提交
338 339

            parent_block.append_op(
340 341
                type="fill_constant_batch_size_like",
                inputs={'Input': [batch_ref]},
Y
Yu Yang 已提交
342 343 344
                outputs={'Out': [boot_var]},
                attrs={
                    'value': init_value,
345
                    'shape': boot_var.shape,
F
fengjiayi 已提交
346
                    'dtype': boot_var.dtype,
347 348
                    'input_dim_idx': ref_batch_dim_idx,
                    'output_dim_idx': init_batch_dim_idx
Y
Yu Yang 已提交
349 350 351 352 353 354
                })

            return self.memory(init=boot_var)
        else:
            pre_mem = self.helper.create_variable(
                name=unique_name("@".join([self.helper.name, "mem"])),
F
fengjiayi 已提交
355
                dtype=init.dtype,
Y
Yu Yang 已提交
356 357 358 359 360 361 362 363 364 365
                shape=init.shape)
            self.memories[pre_mem.name] = StaticRNNMemoryLink(
                init=init, pre_mem=pre_mem)
            return pre_mem

    def step_input(self, x):
        self._assert_in_rnn_block_('step_input')
        if not isinstance(x, Variable):
            raise TypeError("step input takes a Variable")
        if self.seq_len is None:
Y
Yu Yang 已提交
366 367
            self.seq_len = x.shape[0]
        elif self.seq_len != x.shape[0]:
Y
Yu Yang 已提交
368 369 370
            raise ValueError("Static RNN only take fix seq_len input")

        ipt = self.helper.create_variable(
F
fengjiayi 已提交
371
            name=x.name, dtype=x.dtype, shape=list(x.shape[1:]), type=x.type)
Y
Yu Yang 已提交
372 373 374 375 376 377 378 379
        self.inputs.append(ipt)
        return ipt

    def step_output(self, o):
        self._assert_in_rnn_block_('step_output')
        if not isinstance(o, Variable):
            raise TypeError("step output takes a Variable")

F
fengjiayi 已提交
380
        tmp_o = self.helper.create_tmp_variable(dtype=o.dtype)
Y
Yu Yang 已提交
381 382 383 384
        self.helper.append_op(
            type='rnn_memory_helper',
            inputs={'X': [o]},
            outputs={'Out': tmp_o},
F
fengjiayi 已提交
385
            attrs={'dtype': o.dtype})
Y
Yu Yang 已提交
386

Y
Yu Yang 已提交
387
        out_var = self.parent_block().create_var(
Y
Yu Yang 已提交
388 389
            name=tmp_o.name,
            shape=[self.seq_len] + list(tmp_o.shape),
F
fengjiayi 已提交
390
            dtype=tmp_o.dtype)
Y
Yu Yang 已提交
391 392 393 394 395 396 397 398 399 400 401 402 403

        self.outputs.append(out_var)

    def output(self, *outputs):
        for each in outputs:
            self.step_output(each)

    def update_memory(self, mem, var):
        if not isinstance(mem, Variable) or not isinstance(var, Variable):
            raise TypeError("update memory should take variables")
        self.memories[mem.name].mem = var

    def parent_block(self):
404
        prog = self.helper.main_program
Y
Yu Yang 已提交
405 406 407 408 409 410 411 412 413 414 415 416 417 418 419
        parent_idx = prog.current_block().parent_idx
        assert parent_idx >= 0
        parent_block = prog.block(parent_idx)
        return parent_block

    def __call__(self, *args, **kwargs):
        if self.status != StaticRNN.AFTER_RNN_BLOCK:
            raise ValueError("RNN output can only be retrieved after rnn block")
        if len(self.outputs) == 0:
            raise ValueError("RNN has no output")
        elif len(self.outputs) == 1:
            return self.outputs[0]
        else:
            return self.outputs

Y
Yang Yang 已提交
420
    def complete_op(self):
421 422
        main_program = self.helper.main_program
        rnn_block = main_program.current_block()
Y
Yu Yang 已提交
423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461
        parent_block = self.parent_block()

        local_inputs = set()

        for op in rnn_block.ops:
            assert isinstance(op, Operator)
            for oname in op.output_names:
                for out_var_name in op.output(oname):
                    local_inputs.add(out_var_name)

        for var in self.inputs:
            local_inputs.add(var.name)
        for m in self.memories:
            local_inputs.add(m)

        params = list()
        for op in rnn_block.ops:
            assert isinstance(op, Operator)
            for iname in op.input_names:
                for in_var_name in op.input(iname):
                    if in_var_name not in local_inputs:
                        params.append(in_var_name)

        parameters = [parent_block.var(name) for name in params]

        step_scope = parent_block.create_var(
            type=core.VarDesc.VarType.STEP_SCOPES)

        inlinks = [parent_block.var(i.name) for i in self.inputs]
        outlinks = self.outputs

        boot_memories = []
        pre_memories = []
        memories = []
        for _, mem in self.memories.iteritems():
            boot_memories.append(mem.init)
            pre_memories.append(mem.pre_mem.name)
            mem_var = rnn_block.var(mem.mem.name)
            assert isinstance(mem_var, Variable)
F
fengjiayi 已提交
462
            new_mem = self.helper.create_tmp_variable(dtype=mem_var.dtype)
Y
Yu Yang 已提交
463 464 465 466 467

            rnn_block.append_op(
                type='rnn_memory_helper',
                inputs={'X': [mem_var]},
                outputs={'Out': [new_mem]},
F
fengjiayi 已提交
468
                attrs={'dtype': mem_var.dtype})
Y
Yu Yang 已提交
469 470 471 472 473 474 475 476 477 478 479 480 481 482 483

            memories.append(new_mem.name)

        parent_block.append_op(
            type='recurrent',
            inputs={
                'inputs': inlinks,
                'initial_states': boot_memories,
                'parameters': parameters
            },
            outputs={'outputs': outlinks,
                     'step_scopes': [step_scope]},
            attrs={
                'ex_states': pre_memories,
                'states': memories,
484
                'sub_block': rnn_block
Y
Yu Yang 已提交
485
            })
Y
Yu Yang 已提交
486 487


Y
Yang Yang(Tony) 已提交
488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511
class WhileGuard(BlockGuard):
    def __init__(self, while_op):
        if not isinstance(while_op, While):
            raise TypeError("WhileGuard takes a while op")
        super(WhileGuard, self).__init__(while_op.helper.main_program)
        self.while_op = while_op

    def __enter__(self):
        self.while_op.status = While.IN_WHILE_BLOCK
        return super(WhileGuard, self).__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            return False
        self.while_op.status = While.AFTER_WHILE_BLOCK
        self.while_op.complete()
        return super(WhileGuard, self).__exit__(exc_type, exc_val, exc_tb)


class While(object):
    BEFORE_WHILE_BLOCK = 0
    IN_WHILE_BLOCK = 1
    AFTER_WHILE_BLOCK = 2

512 513
    def __init__(self, cond, name=None):
        self.helper = LayerHelper("while", name=name)
Y
Yang Yang(Tony) 已提交
514 515 516 517
        self.status = While.BEFORE_WHILE_BLOCK
        if not isinstance(cond, Variable):
            raise TypeError("condition should be a variable")
        assert isinstance(cond, Variable)
F
fengjiayi 已提交
518
        if cond.dtype != core.DataType.BOOL:
Y
Yang Yang(Tony) 已提交
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560
            raise TypeError("condition should be a bool variable")
        if reduce(lambda a, b: a * b, cond.shape, 1) != 1:
            raise TypeError("condition should be a bool scalar")
        self.cond_var = cond

    def block(self):
        return WhileGuard(self)

    def complete(self):
        main_program = self.helper.main_program
        while_block = main_program.current_block()
        parent_block = main_program.block(main_program.current_block()
                                          .parent_idx)

        inner_outputs = {self.cond_var.name}
        x_name_list = set()
        for op in while_block.ops:
            for iname in op.input_names:
                for in_var_name in op.input(iname):
                    if in_var_name not in inner_outputs:
                        x_name_list.add(in_var_name)

            for oname in op.output_names:
                for out_var_name in op.output(oname):
                    inner_outputs.add(out_var_name)

        out_vars = []
        for inner_out_name in inner_outputs:
            if inner_out_name in parent_block.vars:
                out_vars.append(parent_block.var(inner_out_name))

        step_scope = parent_block.create_var(
            type=core.VarDesc.VarType.STEP_SCOPES)

        parent_block.append_op(
            type='while',
            inputs={
                'X': [parent_block.var(x_name) for x_name in x_name_list],
                'Condition': [self.cond_var]
            },
            outputs={'Out': out_vars,
                     'StepScopes': [step_scope]},
561
            attrs={'sub_block': while_block})
Y
Yang Yang(Tony) 已提交
562 563


564
def lod_rank_table(x, level=0):
Y
yangyaming 已提交
565 566 567
    """LoD Rank Table Operator. Given an input variable **x** and a level number
    of LoD, this layer creates a LodRankTable object. A LoDRankTable object
    contains a list of bi-element tuples. Each tuple consists of an index and
568
    a length, both of which are int type. Refering to specified level of LoD,
Y
yangyaming 已提交
569 570 571
    the index is the sequence index number and the length representes the
    sequence length. Please note that the list is ranked in descending order by
    the length. The following is an example:
Y
yangyaming 已提交
572 573 574 575

        .. code-block:: text

            x is a LoDTensor:
Y
yangyaming 已提交
576
                x.lod = [[0,                2, 3],
Y
yangyaming 已提交
577 578 579
                         [0,             5, 6, 7]]
                x.data = [a, b, c, d, e, f, g]

Y
yangyaming 已提交
580 581 582
            1. set level to 0:
                Create lod rank table:
                    lod_rank_table_obj = lod_rank_table(x, level=0)
Y
yangyaming 已提交
583

Y
yangyaming 已提交
584 585 586 587 588 589 590 591 592
                Get:
                    lod_rank_table_obj.items() = [(0, 2), (1, 1)]

            2. set level to 1:
                Create lod rank table:
                    lod_rank_table_obj = lod_rank_table(x, level=1)

                Get:
                    lod_rank_table_obj.items() = [(0, 5), (1, 1), (2, 1)]
Y
yangyaming 已提交
593 594 595 596

    Args:
        x (Variable): Input variable, a LoDTensor based which to create the lod
            rank table.
Y
yangyaming 已提交
597 598
        level (int): Specify the LoD level, on which to create the lod rank
            table.
Y
yangyaming 已提交
599 600 601 602 603 604 605 606 607 608

    Returns:
        Variable: The created LoDRankTable object.

    Examples:
        .. code-block:: python

            x = fluid.layers.data(name='x', shape=[10],
                            dtype='float32', lod_level=1)
            out = layers.lod_rank_table(x=x, level=0)
609
    """
Y
Yu Yang 已提交
610 611 612 613 614 615 616 617 618 619
    helper = LayerHelper("lod_rank_table", **locals())
    table = helper.create_variable(
        type=core.VarDesc.VarType.LOD_RANK_TABLE,
        name=unique_name("lod_rank_table"))
    helper.append_op(
        type='lod_rank_table',
        inputs={'X': x},
        outputs={'Out': table},
        attrs={'level': level})
    return table
Y
Yu Yang 已提交
620 621


622
def max_sequence_len(rank_table):
Y
yangyaming 已提交
623
    """Max Sequence Len Operator. Given a LoDRankTable object, this layer
Y
yangyaming 已提交
624 625 626 627
    returns the max length of a batch of sequences. In fact, a LoDRankTable
    object contains a list of tuples(<sequence index, sequence length>) and
    the list is already sorted by sequence length in descending order, so the
    operator just returns the sequence length of the first tuple element.
Y
yangyaming 已提交
628 629 630 631 632

    Args:
        rank_table (Variable): Input variable which is a LoDRankTable object.

    Returns:
Y
yangyaming 已提交
633
        Variable: The max length of sequence.
Y
yangyaming 已提交
634 635 636 637 638 639 640 641

    Examples:
        .. code-block:: python

            x = fluid.layers.data(name='x', shape=[10],
                            dtype='float32', lod_level=1)
            rank_table = layers.lod_rank_table(x=x, level=0)
            max_seq_len = layers.max_sequence_len(rank_table)
F
fengjiayi 已提交
642 643 644 645 646 647 648 649 650 651
    """
    helper = LayerHelper("max_seqence_len", **locals())
    res = helper.create_tmp_variable(dtype="int64")
    helper.append_op(
        type="max_sequence_len",
        inputs={"RankTable": rank_table},
        outputs={"Out": res})
    return res


652
def topk(input, k):
653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676
    """
    **topk**

    This function performs the operation that selects the k entries in the input
    vector and outputs their values and indices as vectors. Thus topk_out[j] is
    the j-th largest entry in input, and its index is topk_indices[j]

    Args:
        input (Variable|list): The input tensor that has all the data.
        k (int): The number of top elements that the function will pick.

    Returns:
        Variable: The variable of type array that contains the k largest entries
                  from input.
        Variable: The variable of type array that contains the indices of k
                  largest entries from input.

    Examples:
        .. code-block:: python

          x = fluid.layers.data(name='x', shape=[10])
          k = 5
          array = fluid.layers.topk(x, k)
    """
Y
Yu Yang 已提交
677 678 679 680 681 682 683 684 685 686 687 688
    helper = LayerHelper('topk', **locals())
    topk_out = helper.create_tmp_variable(dtype=input.data_type)
    topk_indices = helper.create_tmp_variable(dtype='int64')
    helper.append_op(
        type='top_k',
        inputs={'X': [input]},
        outputs={'Out': [topk_out],
                 'Indices': [topk_indices]},
        attrs={'k': k})
    return topk_out, topk_indices


689
def lod_tensor_to_array(x, table):
690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708
    """This function performs the operation that converts an LOD_Tensor to
       an array.

    Args:
        x (Variable|list): The tensor that needs to be converted to an array.
        table (ParamAttr|list): The variable that stores the level of lod
                                which is ordered by sequence length in
                                descending order.

    Returns:
        Variable: The variable of type array that has been converted from a
                  tensor.

    Examples:
        .. code-block:: python

          x = fluid.layers.data(name='x', shape=[10])
          table = fluid.layers.lod_rank_table(x, level=0)
          array = fluid.layers.lod_tensor_to_array(x, table)
709
    """
710 711 712
    helper = LayerHelper("lod_tensor_to_array", **locals())
    array = helper.create_variable(
        name=unique_name("lod_tensor_to_array"),
713
        type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
F
fengjiayi 已提交
714
        dtype=x.dtype)
715 716 717 718 719 720 721 722
    helper.append_op(
        type='lod_tensor_to_array',
        inputs={'X': x,
                'RankTable': table},
        outputs={'Out': array})
    return array


723
def array_to_lod_tensor(x, table):
724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743
    """This function performs the operations that converts an array to
       an LOD_Tensor.

    Args:
        x (Variable|list): The array that needs to be converted to a tensor.
        table (ParamAttr|list): The variable that stores the level of lod
                                which is ordered by sequence length in
                                descending order.

    Returns:
        Variable: The variable of type tensor that has been converted
                  from an array.

    Examples:
        .. code-block:: python

          x = fluid.layers.data(name='x', shape=[10])
          table = fluid.layers.lod_rank_table(x, level=0)
          array = fluid.layers.lod_tensor_to_array(x, table)
          lod_tensor = fluid.layers.array_to_lod_tensor(array, table)
744
    """
745
    helper = LayerHelper("array_to_lod_tensor", **locals())
F
fengjiayi 已提交
746
    tmp = helper.create_tmp_variable(dtype=x.dtype)
747 748 749 750 751 752 753 754
    helper.append_op(
        type="array_to_lod_tensor",
        inputs={'X': x,
                'RankTable': table},
        outputs={'Out': tmp})
    return tmp


755
def increment(x, value=1.0, in_place=True):
756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773
    """This function performs an operation that increments each value in the
    input :math:`x` by an amount: :math:`value` as mentioned in the input
    parameter. This operation is performed in-place by default.

    Args:
        x (Variable|list): The tensor that has the input values.
        value (float): The amount by which the values should be incremented.
        in_place (bool): If the increment should be performed in-place.

    Returns:
        Variable: The tensor variable storing the transformation of
                  element-wise increment of each value in the input.

    Examples:
        .. code-block:: python

          data = fluid.layers.data(name='data', shape=[32, 32], dtype='float32')
          data = fluid.layers.increment(x=data, value=3.0, in_place=True)
774
    """
Y
Yu Yang 已提交
775
    helper = LayerHelper("increment", **locals())
Y
Yang Yang(Tony) 已提交
776
    if not in_place:
F
fengjiayi 已提交
777
        out = helper.create_tmp_variable(dtype=x.dtype)
Y
Yang Yang(Tony) 已提交
778 779
    else:
        out = x
Y
Yu Yang 已提交
780 781 782
    helper.append_op(
        type='increment',
        inputs={'X': [x]},
Y
Yang Yu 已提交
783
        outputs={'Out': [out]},
784
        attrs={'step': float(value)})
Y
Yang Yu 已提交
785
    return out
Y
Yu Yang 已提交
786 787


788
def array_write(x, i, array=None):
789
    """This function performs the operation to write the data out as an
790
    LOD_TENSOR_ARRAY.
791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806

    Args:
        x (Variable|list): The input tensor from which the data will be read.
        i (Variable|list): The subscript index in tensor array, that points the
                           place from which data will be read.
        array (Variable|list): The data can be read into this variable if
                               this is assigned.
    Returns:
        Variable: The tensor type variable that has the data written to it.

    Examples:
        .. code-block::python

          tmp = fluid.layers.zeros(shape=[10], dtype='int32')
          i = fluid.layers.fill_constant(shape=[1], dtype='int64', value=10)
          arr = layers.array_write(tmp, i=i)
807
    """
Y
Yu Yang 已提交
808 809 810 811 812
    helper = LayerHelper('array_write', **locals())
    if array is None:
        array = helper.create_variable(
            name="{0}.out".format(helper.name),
            type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
F
fengjiayi 已提交
813
            dtype=x.dtype)
Y
Yu Yang 已提交
814 815 816 817 818 819 820 821
    helper.append_op(
        type='write_to_array',
        inputs={'X': [x],
                'I': [i]},
        outputs={'Out': [array]})
    return array


822
def create_array(dtype):
823 824 825 826 827 828 829 830 831 832 833 834 835 836 837
    """This function creates an array of type :math:`LOD_TENSOR_ARRAY` using the
    LayerHelper.

    Args:
        dtype (int|float): The data type of the elements in the array.

    Returns:
        Variable: The tensor variable storing the elements of data type.

    Examples:
        .. code-block:: python

          data = fluid.layers.create_array(dtype='float32')

    """
Y
Yang Yang(Tony) 已提交
838 839 840 841 842 843 844
    helper = LayerHelper("array", **locals())
    return helper.create_variable(
        name="{0}.out".format(helper.name),
        type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
        dtype=dtype)


845
def less_than(x, y, cond=None, **ignored):
846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863
    """
    **Less than**

    This layer returns the truth value of :math:`x < y` elementwise.

    Args:
        x(Variable): First operand of *less_than*
        y(Variable): Second operand of *less_than*
        cond(Variable|None): Optional output variable to store the result of *less_than*

    Returns:
        Variable: The tensor variable storing the output of *less_than*.

    Examples:
        .. code-block:: python

          less = fluid.layers.less_than(x=label, y=limit)
    """
Y
Yang Yang(Tony) 已提交
864 865 866 867 868 869 870 871 872 873 874
    helper = LayerHelper("less_than", **locals())
    if cond is None:
        cond = helper.create_tmp_variable(dtype='bool')
        cond.stop_gradient = True

    helper.append_op(
        type='less_than', inputs={'X': [x],
                                  'Y': [y]}, outputs={'Out': [cond]})
    return cond


875
def array_read(array, i):
K
kavyasrinet 已提交
876
    """This function performs the operation to read the data in as an
877
    LOD_TENSOR_ARRAY.
K
kavyasrinet 已提交
878 879 880 881 882 883 884 885 886 887 888
    Args:
        array (Variable|list): The input tensor that will be written to an array.
        i (Variable|list): The subscript index in tensor array, that points the
                           place where data will be written to.
    Returns:
        Variable: The tensor type variable that has the data written to it.
    Examples:
        .. code-block::python
          tmp = fluid.layers.zeros(shape=[10], dtype='int32')
          i = fluid.layers.fill_constant(shape=[1], dtype='int64', value=10)
          arr = layers.array_read(tmp, i=i)
889
    """
Y
Yu Yang 已提交
890 891 892 893 894
    helper = LayerHelper('array_read', **locals())
    if not isinstance(
            array,
            Variable) or array.type != core.VarDesc.VarType.LOD_TENSOR_ARRAY:
        raise TypeError("array should be tensor array vairable")
F
fengjiayi 已提交
895
    out = helper.create_tmp_variable(dtype=array.dtype)
Y
Yu Yang 已提交
896 897 898 899 900 901
    helper.append_op(
        type='read_from_array',
        inputs={'X': [array],
                'I': [i]},
        outputs={'Out': [out]})
    return out
Y
Yang Yu 已提交
902 903


904
def shrink_memory(x, i, table):
905 906 907 908
    """
    This function creates an operator to shrink_rnn_memory using the RankTable
    as mentioned in the input parameter.
    """
Y
Yang Yu 已提交
909
    helper = LayerHelper('shrink_memory', **locals())
F
fengjiayi 已提交
910
    out = helper.create_tmp_variable(dtype=x.dtype)
Y
Yang Yu 已提交
911
    helper.append_op(
Y
Yang Yu 已提交
912
        type='shrink_rnn_memory',
Y
Yang Yu 已提交
913 914 915 916 917 918
        inputs={'X': [x],
                'I': [i],
                'RankTable': [table]},
        outputs={'Out': [out]},
        attrs={})
    return out
Y
Yang Yu 已提交
919 920


921
def array_length(array):
K
kavyasrinet 已提交
922
    """This function performs the operation to find the length of the input
923
    LOD_TENSOR_ARRAY.
K
kavyasrinet 已提交
924 925 926 927 928 929 930 931 932 933 934 935 936 937 938

    Args:
        array (LOD_TENSOR_ARRAY): The input array that will be used
                                  to compute the length.

    Returns:
        Variable: The length of the input LoDTensorArray.

    Examples:
        .. code-block::python

          tmp = fluid.layers.zeros(shape=[10], dtype='int32')
          i = fluid.layers.fill_constant(shape=[1], dtype='int64', value=10)
          arr = fluid.layers.array_write(tmp, i=i)
          arr_len = fluid.layers.array_length(arr)
939
    """
Y
Yang Yu 已提交
940 941 942 943 944 945
    helper = LayerHelper('array_length', **locals())
    tmp = helper.create_tmp_variable(dtype='int64')
    tmp.stop_gradient = True
    helper.append_op(
        type='lod_array_length', inputs={'X': [array]}, outputs={'Out': [tmp]})
    return tmp
Y
Yu Yang 已提交
946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964


class ConditionalBlockGuard(BlockGuard):
    def __init__(self, block):
        if not isinstance(block, ConditionalBlock):
            raise TypeError("block should be conditional block")
        super(ConditionalBlockGuard, self).__init__(block.helper.main_program)
        self.block = block

    def __enter__(self):
        return super(ConditionalBlockGuard, self).__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.block.complete()
        return super(ConditionalBlockGuard, self).__exit__(exc_type, exc_val,
                                                           exc_tb)


class ConditionalBlock(object):
965
    def __init__(self, inputs, name=None):
Y
Yu Yang 已提交
966 967 968 969
        for each_input in inputs:
            if not isinstance(each_input, Variable):
                raise TypeError("Each input should be variable")
        self.inputs = inputs
970
        self.helper = LayerHelper('conditional_block', name=name)
Y
Yu Yang 已提交
971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000

    def block(self):
        return ConditionalBlockGuard(self)

    def complete(self):
        inside_block = self.helper.main_program.current_block()
        parent_block = self.helper.main_program.block(inside_block.parent_idx)

        intermediate = set()
        params = set()

        for each_op in inside_block.ops:
            assert isinstance(each_op, Operator)
            for iname in each_op.input_names:
                for in_var_name in each_op.input(iname):
                    if in_var_name not in intermediate:
                        params.add(in_var_name)

            for oname in each_op.output_names:
                for out_var_name in each_op.output(oname):
                    intermediate.add(out_var_name)
        input_set = set([ipt.name for ipt in self.inputs])

        param_list = [
            parent_block.var(each_name) for each_name in params
            if each_name not in input_set
        ]

        out_list = [
            parent_block.var(var_name) for var_name in parent_block.vars
X
xuwei06 已提交
1001
            if var_name in intermediate
Y
Yu Yang 已提交
1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013
        ]

        step_scope = parent_block.create_var(
            type=core.VarDesc.VarType.STEP_SCOPES)
        parent_block.append_op(
            type='conditional_block',
            inputs={
                'X': self.inputs,
                'Params': param_list,
            },
            outputs={'Out': out_list,
                     'Scope': [step_scope]},
1014
            attrs={'sub_block': inside_block})
Y
Yu Yang 已提交
1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054


class IfElseBlockGuard(object):
    def __init__(self, is_true, ifelse):
        if not isinstance(ifelse, IfElse):
            raise TypeError("ifelse must be an instance of IfElse class")

        if ifelse.status != IfElse.OUT_IF_ELSE_BLOCKS:
            raise ValueError("You cannot invoke IfElse.block() inside a block")

        self.is_true = is_true
        self.ie = ifelse
        if is_true:
            self.cond_block = ifelse.conditional_true_block
        else:
            self.cond_block = ifelse.conditional_false_block

        if not isinstance(self.cond_block, ConditionalBlock):
            raise TypeError("Unexpected situation")

        self.cond_block = self.cond_block.block()

    def __enter__(self):
        self.ie.status = IfElse.IN_IF_ELSE_TRUE_BLOCKS if self.is_true else IfElse.IN_IF_ELSE_FALSE_BLOCKS
        self.cond_block.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self.cond_block.__exit__(exc_type, exc_val, exc_tb):
            # re-raise inside exception
            return False
        if len(self.ie.output_table[1 if self.is_true else 0]) == 0:
            raise ValueError("Must set output inside block")
        self.ie.status = IfElse.OUT_IF_ELSE_BLOCKS


class IfElse(object):
    OUT_IF_ELSE_BLOCKS = 0
    IN_IF_ELSE_TRUE_BLOCKS = 1
    IN_IF_ELSE_FALSE_BLOCKS = 2

1055
    def __init__(self, cond, name=None):
Y
Yu Yang 已提交
1056 1057
        if not isinstance(cond, Variable):
            raise TypeError("cond must be a Variable")
1058
        self.helper = LayerHelper('ifelse', name=name)
Y
Yu Yang 已提交
1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072
        self.cond = cond
        self.input_table = {}
        self.status = IfElse.OUT_IF_ELSE_BLOCKS
        self.conditional_true_block = ConditionalBlock(inputs=[self.cond])
        self.conditional_false_block = ConditionalBlock(inputs=[self.cond])
        self.output_table = ([], [])  # (true_outs, false_outs)

    def input(self, x):
        if self.status == IfElse.OUT_IF_ELSE_BLOCKS:
            raise ValueError("input must in true/false blocks")
        if id(x) not in self.input_table:
            parent_block = self.parent_block()
            out_true = parent_block.create_var(
                name=unique_name('ifelse_input' + self.helper.name),
F
fengjiayi 已提交
1073
                dtype=x.dtype)
Y
Yu Yang 已提交
1074 1075 1076

            out_false = parent_block.create_var(
                name=unique_name('ifelse_input' + self.helper.name),
F
fengjiayi 已提交
1077
                dtype=x.dtype)
Y
Yu Yang 已提交
1078 1079 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 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118
            parent_block.append_op(
                type='split_lod_tensor',
                inputs={
                    'X': x,
                    'Mask': self.cond,
                },
                outputs={'OutTrue': out_true,
                         'OutFalse': out_false},
                attrs={'level': 0})
            self.input_table[id(x)] = (out_true, out_false)
        else:
            out_true, out_false = self.input_table[id(x)]

        if self.status == IfElse.IN_IF_ELSE_TRUE_BLOCKS:
            return out_true
        else:
            return out_false

    def parent_block(self):
        current_block = self.helper.main_program.current_block()
        return self.helper.main_program.block(current_block.parent_idx)

    def true_block(self):
        return IfElseBlockGuard(True, self)

    def false_block(self):
        return IfElseBlockGuard(False, self)

    def output(self, *outs):
        if self.status == self.OUT_IF_ELSE_BLOCKS:
            raise ValueError("output can only be invoked in the sub-block")

        out_table = self.output_table[1 if self.status ==
                                      self.IN_IF_ELSE_TRUE_BLOCKS else 0]
        parent_block = self.parent_block()
        for each_out in outs:
            if not isinstance(each_out, Variable):
                raise TypeError("Each output should be a variable")
            # create outside tensor
            outside_out = parent_block.create_var(
                name=unique_name("_".join([self.helper.name, 'output'])),
F
fengjiayi 已提交
1119
                dtype=each_out.dtype)
Y
Yu Yang 已提交
1120 1121 1122
            out_table.append(outside_out)

            # assign local var to outside
1123
            assign(input=each_out, output=outside_out)
Y
Yu Yang 已提交
1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146

    def __call__(self):
        if self.status != self.OUT_IF_ELSE_BLOCKS:
            raise ValueError("IfElse::__call__ must be out of sub-block")
        false_len, true_len = map(len, self.output_table)
        if false_len == 0 and true_len == 0:
            raise ValueError("Must invoke true_block/false_block before "
                             "__call__")
        elif false_len != true_len and false_len != 0 and true_len != 0:
            raise ValueError("The output side must be same")
        elif false_len == 0 or true_len == 0:
            return self.output_table[0 if false_len != 0 else 1]

        # else none of false_len/true_len is zero
        # merge together
        rlist = []
        for false_var, true_var in zip(*self.output_table):
            rlist.append(
                merge_lod_tensor(
                    in_true=true_var,
                    in_false=false_var,
                    mask=self.cond,
                    x=self.cond,
1147
                    level=0))
Y
Yu Yang 已提交
1148
        return rlist
1149 1150 1151 1152 1153 1154 1155


class DynamicRNN(object):
    BEFORE_RNN = 0
    IN_RNN = 1
    AFTER_RNN = 2

1156 1157
    def __init__(self, name=None):
        self.helper = LayerHelper('dynamic_rnn', name=name)
1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210
        self.status = DynamicRNN.BEFORE_RNN
        self.lod_rank_table = None
        self.max_seq_len = None
        self.step_idx = None
        self.zero_idx = fill_constant(shape=[1], value=0, dtype='int64')
        self.mem_dict = dict()
        self.output_array = []
        self.outputs = []
        self.cond = self.helper.create_tmp_variable(dtype='bool')
        self.cond.stop_gradient = False
        self.while_op = While(self.cond)
        self.input_array = []
        self.mem_link = []

    def step_input(self, x):
        self._assert_in_rnn_block_("step_input")
        if not isinstance(x, Variable):
            raise TypeError(
                "step_input() can only take a Variable as its input")
        parent_block = self._parent_block_()
        if self.lod_rank_table is None:
            self.lod_rank_table = parent_block.create_var(
                name=unique_name('lod_rank_table'),
                type=core.VarDesc.VarType.LOD_RANK_TABLE)
            self.lod_rank_table.stop_gradient = True
            parent_block.append_op(
                type='lod_rank_table',
                inputs={"X": x},
                outputs={"Out": self.lod_rank_table})
            self.max_seq_len = parent_block.create_var(
                name=unique_name('dynamic_rnn_max_seq_len'), dtype='int64')
            self.max_seq_len.stop_gradient = False
            parent_block.append_op(
                type='max_sequence_len',
                inputs={'RankTable': self.lod_rank_table},
                outputs={"Out": self.max_seq_len})
            self.cond.stop_gradient = True
            parent_block.append_op(
                type='less_than',
                inputs={'X': self.step_idx,
                        'Y': self.max_seq_len},
                outputs={'Out': self.cond})

        input_array = parent_block.create_var(
            name=unique_name('dynamic_rnn_input_array'),
            type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
            dtype=x.dtype)
        self.input_array.append((input_array, x.dtype))
        parent_block.append_op(
            type='lod_tensor_to_array',
            inputs={'X': x,
                    'RankTable': self.lod_rank_table},
            outputs={'Out': input_array})
1211
        return array_read(array=input_array, i=self.step_idx)
1212 1213 1214 1215 1216 1217 1218 1219 1220 1221

    @contextlib.contextmanager
    def block(self):
        if self.status != DynamicRNN.BEFORE_RNN:
            raise ValueError("rnn.block() can only be invoke once")
        self.step_idx = fill_constant(shape=[1], dtype='int64', value=0)
        self.step_idx.stop_gradient = False
        self.status = DynamicRNN.IN_RNN
        with self.while_op.block():
            yield
1222
            increment(x=self.step_idx, value=1.0, in_place=True)
1223 1224

            for new_mem, mem_array in self.mem_link:
1225 1226 1227
                array_write(x=new_mem, i=self.step_idx, array=mem_array)

            less_than(x=self.step_idx, y=self.max_seq_len, cond=self.cond)
1228 1229 1230 1231 1232

        self.status = DynamicRNN.AFTER_RNN
        for each_array in self.output_array:
            self.outputs.append(
                array_to_lod_tensor(
1233
                    x=each_array, table=self.lod_rank_table))
1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259

    def __call__(self, *args, **kwargs):
        if self.status != DynamicRNN.AFTER_RNN:
            raise ValueError(
                "Dynamic RNN outputs can only be retrieved after rnn block")
        if len(self.outputs) == 1:
            return self.outputs[0]
        else:
            return self.outputs

    def memory(self, init=None, shape=None, value=0.0, dtype='float32'):
        self._assert_in_rnn_block_('memory')
        if init is not None:
            if not isinstance(init, Variable):
                raise TypeError(
                    "The input arg `init` of memory() must be a Variable")
            parent_block = self._parent_block_()
            mem_array = parent_block.create_var(
                name=unique_name('dynamic_rnn_mem_array'),
                type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
                dtype=init.dtype)
            parent_block.append_op(
                type='write_to_array',
                inputs={'X': init,
                        'I': self.zero_idx},
                outputs={'Out': mem_array})
1260
            retv = array_read(array=mem_array, i=self.step_idx)
1261
            retv = shrink_memory(
1262
                x=retv, i=self.step_idx, table=self.lod_rank_table)
1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331
            self.mem_dict[retv.name] = mem_array
            return retv
        else:
            if len(self.input_array) == 0:
                raise ValueError(
                    "step_input should be invoked before memory(shape=..., value=...)"
                )
            parent_block = self._parent_block_()
            init = parent_block.create_var(
                name=unique_name('mem_init'), dtype=dtype)
            arr, dtype = self.input_array[0]
            in0 = parent_block.create_var(name=unique_name('in0'), dtype=dtype)
            parent_block.append_op(
                type='read_from_array',
                inputs={'X': [arr],
                        'I': [self.zero_idx]},
                outputs={'Out': [in0]})
            parent_block.append_op(
                type='fill_constant_batch_size_like',
                inputs={'Input': [in0]},
                outputs={'Out': [init]},
                attrs={
                    'shape': [-1] + shape,
                    'value': float(value),
                    'dtype': init.dtype
                })
            return self.memory(init=init)

    def update_memory(self, ex_mem, new_mem):
        self._assert_in_rnn_block_('update_memory')
        if not isinstance(ex_mem, Variable):
            raise TypeError("The input arg `ex_mem` of update_memory() must "
                            "be a Variable")
        if not isinstance(new_mem, Variable):
            raise TypeError("The input arg `new_mem` of update_memory() must "
                            "be a Variable")

        mem_array = self.mem_dict.get(ex_mem.name, None)
        if mem_array is None:
            raise ValueError("Please invoke memory before update_memory")
        if self.lod_rank_table is None:
            raise ValueError("Please invoke step_input before update_memory")

        self.mem_link.append((new_mem, mem_array))

    def output(self, *outputs):
        self._assert_in_rnn_block_('output')
        parent_block = self._parent_block_()
        for each in outputs:
            outside_array = parent_block.create_var(
                name=unique_name("_".join(
                    [self.helper.name, "output_array", each.name])),
                type=core.VarDesc.VarType.LOD_TENSOR_ARRAY,
                dtype=each.dtype)
            array_write(x=each, i=self.step_idx, array=outside_array)
            self.output_array.append(outside_array)

    def _parent_block_(self):
        prog = self.helper.main_program
        parent_idx = prog.current_block().parent_idx
        assert parent_idx >= 0
        parent_block = prog.block(parent_idx)

        return parent_block

    def _assert_in_rnn_block_(self, method):
        if self.status != DynamicRNN.IN_RNN:
            raise ValueError("{0} can only be invoked inside rnn block.".format(
                method))
Y
Yang Yu 已提交
1332 1333


Y
Yang Yu 已提交
1334
@autodoc
Y
Yang Yu 已提交
1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346
def reorder_lod_tensor_by_rank(x, rank_table):
    helper = LayerHelper('reorder_lod_tensor_by_rank', **locals())
    helper.is_instance('x', Variable)
    helper.is_instance('rank_table', Variable)

    out = helper.create_tmp_variable(dtype=x.dtype)
    helper.append_op(
        type='reorder_lod_tensor_by_rank',
        inputs={'X': [x],
                'RankTable': [rank_table]},
        outputs={'Out': [out]})
    return out