Map on the list of tensors unpacked from labeled_tensor. See tf.map_fn. Args: fn: The function to apply to each unpacked LabeledTensor. It should have type LabeledTensor -> LabeledTensor. labeled_tensor: The input tensor. name: Optional op name. Returns: A tensor that
(fn, labeled_tensor, name=None)
| 609 | @tc.accepts(collections_abc.Callable, core.LabeledTensorLike, |
| 610 | tc.Optional(string_types)) |
| 611 | def map_fn(fn, labeled_tensor, name=None): |
| 612 | """Map on the list of tensors unpacked from labeled_tensor. |
| 613 | |
| 614 | See tf.map_fn. |
| 615 | |
| 616 | Args: |
| 617 | fn: The function to apply to each unpacked LabeledTensor. |
| 618 | It should have type LabeledTensor -> LabeledTensor. |
| 619 | labeled_tensor: The input tensor. |
| 620 | name: Optional op name. |
| 621 | |
| 622 | Returns: |
| 623 | A tensor that packs the results of applying fn to the list of tensors |
| 624 | unpacked from labeled_tensor. |
| 625 | """ |
| 626 | with ops.name_scope(name, 'lt_map_fn', [labeled_tensor]) as scope: |
| 627 | labeled_tensor = core.convert_to_labeled_tensor(labeled_tensor) |
| 628 | |
| 629 | unpack_lts = unpack(labeled_tensor) |
| 630 | |
| 631 | # TODO(ericmc): Fix this upstream. |
| 632 | if labeled_tensor.dtype == dtypes.string: |
| 633 | # We must construct the full graph here, because map_fn_lib.map_fn |
| 634 | # doesn't work for string-valued tensors. |
| 635 | # Constructing the full graph may be slow. |
| 636 | map_lts = [fn(t) for t in unpack_lts] |
| 637 | return pack(map_lts, list(labeled_tensor.axes.values())[0], name=scope) |
| 638 | else: |
| 639 | # Figure out what the axis labels should be, but use tf.map_fn to |
| 640 | # construct the graph because it's efficient. |
| 641 | # It may be slow to construct the full graph, so we infer the labels from |
| 642 | # the first element. |
| 643 | # TODO(ericmc): This builds a subgraph which then gets thrown away. |
| 644 | # Find a more elegant solution. |
| 645 | first_map_lt = fn(unpack_lts[0]) |
| 646 | final_axes = list(labeled_tensor.axes.values())[:1] + list( |
| 647 | first_map_lt.axes.values()) |
| 648 | |
| 649 | @tc.returns(ops.Tensor) |
| 650 | @tc.accepts(ops.Tensor) |
| 651 | def tf_fn(tensor): |
| 652 | original_axes = list(labeled_tensor.axes.values())[1:] |
| 653 | tensor_lt = core.LabeledTensor(tensor, original_axes) |
| 654 | return fn(tensor_lt).tensor |
| 655 | |
| 656 | map_op = map_fn_lib.map_fn( |
| 657 | tf_fn, labeled_tensor.tensor, dtype=first_map_lt.dtype) |
| 658 | map_lt = core.LabeledTensor(map_op, final_axes) |
| 659 | |
| 660 | return core.identity(map_lt, name=scope) |
| 661 | |
| 662 | |
| 663 | @tc.returns(core.LabeledTensor) |