An operand for a tflite hint function that is aggregated from many. For example, an LSTM is a grid of operators that are all related. Inputs going into them may need to be fused, so they should all be tracked as related arguments.
| 536 | |
| 537 | |
| 538 | class _LiteAggregateOperand(_LiteOperand): |
| 539 | """An operand for a tflite hint function that is aggregated from many. |
| 540 | |
| 541 | For example, an LSTM is a grid of operators that are all related. Inputs |
| 542 | going into them may need to be fused, so they should all be tracked as |
| 543 | related arguments. |
| 544 | """ |
| 545 | |
| 546 | def __init__(self, aggregation): |
| 547 | _LiteOperand.__init__(self) |
| 548 | self.aggregation = aggregation |
| 549 | self.names = {} |
| 550 | self.nodes = {} |
| 551 | self.flattened = None |
| 552 | |
| 553 | def add(self, sort, node): |
| 554 | self.names[sort] = _tensor_name_base(node.name) |
| 555 | self.nodes[sort] = node |
| 556 | |
| 557 | def flatten_nodes(self): |
| 558 | """Return a list of all the node protos in aggregation sorted order.""" |
| 559 | if not self.flattened: |
| 560 | self.flattened = [None] * len(self.nodes) |
| 561 | for idx, node in _six.iteritems(self.nodes): |
| 562 | self.flattened[idx] = node |
| 563 | for n in self.nodes: |
| 564 | if n is None: |
| 565 | raise RuntimeError("Aggregate was missing argument.") |
| 566 | if self.aggregation == OpHint.AGGREGATE_FIRST: |
| 567 | self.flattened = self.flattened[:1] |
| 568 | elif self.aggregation == OpHint.AGGREGATE_LAST: |
| 569 | self.flattened = self.flattened[-1:] |
| 570 | elif self.aggregation == OpHint.AGGREGATE_STACK: |
| 571 | pass |
| 572 | else: |
| 573 | raise ValueError( |
| 574 | "Invalid aggregation type %r specified" % self.aggregation) |
| 575 | return self.flattened |
| 576 | |
| 577 | def flatten(self): |
| 578 | """Return a list of all node names in aggregation sorted sorter.""" |
| 579 | return [_tensor_name_base(x.name) for x in self.flatten_nodes()] |
| 580 | |
| 581 | def aggregate_and_return_name_for_input(self, out_graphdef): |
| 582 | """This adds the nodes to out_graphdef and returns an aggregated output. |
| 583 | |
| 584 | In particular, if you have 4 inputs to a hint stub, this will be the |
| 585 | node that you can use as an output. I.e. you have 4 timesteps from a |
| 586 | static rnn, then a fused UnidriecitonalLSTM will expect 1 input with |
| 587 | all 4 time steps. So here we make a pack and return the output name of |
| 588 | that pack. |
| 589 | |
| 590 | Args: |
| 591 | out_graphdef: A graphdef that is ready to have this input added. |
| 592 | |
| 593 | Returns: |
| 594 | The name of a pack that aggregates this node. |
| 595 | """ |