A function wrapper that supports structured arguments and return values.
| 2641 | |
| 2642 | |
| 2643 | class StructuredFunctionWrapper(object): |
| 2644 | """A function wrapper that supports structured arguments and return values.""" |
| 2645 | |
| 2646 | # pylint: disable=protected-access |
| 2647 | def __init__(self, |
| 2648 | func, |
| 2649 | transformation_name, |
| 2650 | dataset=None, |
| 2651 | input_classes=None, |
| 2652 | input_shapes=None, |
| 2653 | input_types=None, |
| 2654 | input_structure=None, |
| 2655 | add_to_graph=True, |
| 2656 | use_legacy_function=False, |
| 2657 | defun_kwargs=None): |
| 2658 | """Creates a new `StructuredFunctionWrapper` for the given function. |
| 2659 | |
| 2660 | Args: |
| 2661 | func: A function from a nested structure to another nested structure. |
| 2662 | transformation_name: Human-readable name of the transformation in which |
| 2663 | this function is being instantiated, for error messages. |
| 2664 | dataset: (Optional.) A `tf.data.Dataset`. If given, the structure of this |
| 2665 | dataset will be assumed as the structure for `func` arguments; otherwise |
| 2666 | `input_classes`, `input_shapes`, and `input_types` must be defined. |
| 2667 | input_classes: (Optional.) A nested structure of `type`. If given, this |
| 2668 | argument defines the Python types for `func` arguments. |
| 2669 | input_shapes: (Optional.) A nested structure of `tf.TensorShape`. If |
| 2670 | given, this argument defines the shapes and structure for `func` |
| 2671 | arguments. |
| 2672 | input_types: (Optional.) A nested structure of `tf.DType`. If given, this |
| 2673 | argument defines the element types and structure for `func` arguments. |
| 2674 | input_structure: (Optional.) A `Structure` object. If given, this argument |
| 2675 | defines the element types and structure for `func` arguments. |
| 2676 | add_to_graph: (Optional.) If `True`, the function will be added to the |
| 2677 | default graph. |
| 2678 | use_legacy_function: (Optional.) A boolean that determines whether the |
| 2679 | function be created using `tensorflow.python.eager.function.defun` |
| 2680 | (default behavior) or `tensorflow.python.framework.function.Defun` |
| 2681 | (legacy beheavior). |
| 2682 | defun_kwargs: (Optional.) A dictionary mapping string argument names to |
| 2683 | values. If supplied, will be passed to `function` as keyword arguments. |
| 2684 | |
| 2685 | Raises: |
| 2686 | ValueError: If an invalid combination of `dataset`, `input_classes`, |
| 2687 | `input_shapes`, and `input_types` is passed. |
| 2688 | """ |
| 2689 | if input_structure is None: |
| 2690 | if dataset is None: |
| 2691 | if input_classes is None or input_shapes is None or input_types is None: |
| 2692 | raise ValueError("Either `dataset`, `input_structure` or all of " |
| 2693 | "`input_classes`, `input_shapes`, and `input_types` " |
| 2694 | "must be specified.") |
| 2695 | self._input_structure = structure.convert_legacy_structure( |
| 2696 | input_types, input_shapes, input_classes) |
| 2697 | else: |
| 2698 | if not (input_classes is None and input_shapes is None and |
| 2699 | input_types is None): |
| 2700 | raise ValueError("Either `dataset`, `input_structure` or all of " |