(pfor_input)
| 3136 | @RegisterPFor("StatefulPartitionedCall") |
| 3137 | @RegisterPFor("PartitionedCall") |
| 3138 | def _convert_partitioned_call(pfor_input): |
| 3139 | func_name = pfor_input.get_attr("f").name |
| 3140 | func = pfor_input.op.graph._get_function(compat.as_bytes(func_name)) |
| 3141 | assert isinstance(func.graph, func_graph.FuncGraph), ( |
| 3142 | "Could not find FuncGraph object for %s. Got func %s" % (func_name, func)) |
| 3143 | pfor = pfor_input.pfor |
| 3144 | converter = PFor(loop_var=pfor.loop_var, |
| 3145 | loop_len=pfor.loop_len_vector[0], |
| 3146 | pfor_ops=func.graph.get_operations(), |
| 3147 | all_indices=pfor.all_indices, |
| 3148 | all_indices_partitioned=pfor.all_indices_partitioned, |
| 3149 | pfor_config=pfor.pfor_config) |
| 3150 | |
| 3151 | # TODO(agarwal): consider caching this function definition. |
| 3152 | @def_function.function |
| 3153 | def f(*args): |
| 3154 | assert all(isinstance(arg, WrappedTensor) for arg in args), args |
| 3155 | assert len(args) == len(func.graph.inputs), (args, func.graph.inputs) |
| 3156 | # Map inputs to function arguments. |
| 3157 | for inp, arg in zip(func.graph.inputs, args): |
| 3158 | converter._add_conversion(inp, arg) |
| 3159 | # Convert output tensors. |
| 3160 | return tuple([converter._convert_helper(x).t |
| 3161 | for x in func._func_graph_outputs]) |
| 3162 | |
| 3163 | call_outputs = f(*pfor_input.inputs) |
| 3164 | assert len(call_outputs) == len(func._func_graph_outputs) |
| 3165 | outputs = [] |
| 3166 | for call_output, output_tensor in zip(call_outputs, func._func_graph_outputs): |
| 3167 | func_output = converter._convert_helper(output_tensor) |
| 3168 | outputs.append(wrap(call_output, |
| 3169 | func_output.is_stacked, |
| 3170 | func_output.is_sparse_stacked)) |
| 3171 | return outputs |
nothing calls this directly
no test coverage detected