Maps python function args to graph-construction inputs. Args: args: A flat list of user-specified arguments. names: A list of strings with user-specified argument names, same length as `args`. May be `None`, in which case a generic name is used. structure: The original argument
(args, names, structure, flat_shapes=None)
| 1063 | |
| 1064 | |
| 1065 | def _get_defun_inputs(args, names, structure, flat_shapes=None): |
| 1066 | """Maps python function args to graph-construction inputs. |
| 1067 | |
| 1068 | Args: |
| 1069 | args: A flat list of user-specified arguments. |
| 1070 | names: A list of strings with user-specified argument names, same length as |
| 1071 | `args`. May be `None`, in which case a generic name is used. |
| 1072 | structure: The original argument list or dictionary. |
| 1073 | flat_shapes: A flat list of values that are either `None` or |
| 1074 | instances of `TensorShape`. If provided, then length must match |
| 1075 | that of `nest.flatten(args, expand_composites=True)`; and locations where |
| 1076 | `args` are instances of `Tensor` must have a corresponding `TensorShape` |
| 1077 | in `flat_shapes`. May be `None`, in which case exact shapes are read |
| 1078 | directly from the args. |
| 1079 | |
| 1080 | Returns: |
| 1081 | Placeholders with the same structure as `structure`. |
| 1082 | |
| 1083 | Raises: |
| 1084 | RuntimeError: if `flat_shapes` is provided, but |
| 1085 | `len(flat_shapes) != len(nest.flatten(args, expand_composites=True))`. |
| 1086 | RuntimeError: if a shape from `flat_shapes` is not None |
| 1087 | for an argument that is not a `Tensor`, `TensorSpec`, |
| 1088 | or `ResourceVariable`. |
| 1089 | """ |
| 1090 | func_graph = ops.get_default_graph() |
| 1091 | function_inputs = [] |
| 1092 | if names is None: |
| 1093 | names = [None] * len(args) |
| 1094 | if flat_shapes is None: |
| 1095 | shapes_iter = itertools.repeat(None) |
| 1096 | else: |
| 1097 | len_flat_args = len(nest.flatten(args, expand_composites=True)) |
| 1098 | if len_flat_args != len(flat_shapes): |
| 1099 | raise RuntimeError( |
| 1100 | "Length of fully flat shapes (%d) must match that of " |
| 1101 | "flatten(args) (%d). args: %s, flat_shapes: %s" |
| 1102 | % (len(flat_shapes), |
| 1103 | len_flat_args, |
| 1104 | args, |
| 1105 | flat_shapes)) |
| 1106 | shapes_iter = iter(flat_shapes) |
| 1107 | for arg_value, name in zip(args, names): |
| 1108 | flattened = nest.flatten(arg_value, expand_composites=True) |
| 1109 | tensor_specs = [ |
| 1110 | arg for arg in flattened if isinstance(arg, tensor_spec.TensorSpec) |
| 1111 | ] |
| 1112 | specified_names = [arg.name for arg in tensor_specs if arg.name] |
| 1113 | if specified_names and len(specified_names) < len(tensor_specs): |
| 1114 | raise ValueError("If specifying TensorSpec names for nested structures, " |
| 1115 | "either zero or all names have to be specified.") |
| 1116 | |
| 1117 | for arg in flattened: |
| 1118 | # We have a shape entry for each arg, regadless of whether it's a real |
| 1119 | # Tensor or not. For non-tensor entries it should be None. |
| 1120 | shape = next(shapes_iter) |
| 1121 | if isinstance(arg, (ops.Tensor, tensor_spec.TensorSpec)): |
| 1122 | if isinstance(arg, tensor_spec.TensorSpec) and arg.name: |