Add NextIteration and back edge from v to m.
(m, v, enforce_shape_invariant=True)
| 589 | |
| 590 | |
| 591 | def _AddNextAndBackEdge(m, v, enforce_shape_invariant=True): |
| 592 | """Add NextIteration and back edge from v to m.""" |
| 593 | if isinstance(m, ops.Tensor): |
| 594 | v = ops.convert_to_tensor(v) |
| 595 | v = _NextIteration(v) |
| 596 | if enforce_shape_invariant: |
| 597 | # Make sure the shapes of loop outputs are correct. We do this before |
| 598 | # calling _update_input, which will raise a less-helpful error message if |
| 599 | # the types don't match. |
| 600 | # TODO(skyewm): call this for other cases below (needs testing) |
| 601 | _EnforceShapeInvariant(m, v) |
| 602 | m.op._update_input(1, v) # pylint: disable=protected-access |
| 603 | elif isinstance(m, composite_tensor.CompositeTensor): |
| 604 | # pylint: disable=protected-access |
| 605 | def update_component(m_component, v_component): |
| 606 | m_component.op._update_input(1, v_component) |
| 607 | |
| 608 | if isinstance(m, ops.IndexedSlices): |
| 609 | v = math_ops._as_indexed_slices(v, optimize=False) |
| 610 | # pylint: enable=protected-access |
| 611 | v = _NextIteration(v) |
| 612 | return nest.map_structure(update_component, m, v, expand_composites=True) |
| 613 | else: |
| 614 | raise TypeError("Type %s not supported" % type(m)) |
| 615 | return v |
| 616 | |
| 617 | |
| 618 | @six.add_metaclass(abc.ABCMeta) |
no test coverage detected