Searched refs:retval_shapes (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.h | 296 std::vector<InferredShape>* retval_shapes);
309 const std::vector<InferredShape>& retval_shapes, const Graph& graph,
|
| D | distributed_tpu_rewrite_pass.cc | 1695 std::vector<InferredShape>* retval_shapes) { in GetArgAndRetvalShapes() argument
1797 retval_shapes->clear(); in GetArgAndRetvalShapes()
1800 retval_shapes->resize(node.num_outputs()); in GetArgAndRetvalShapes()
1802 (*retval_shapes)[i].shape = it->second[i].shape; in GetArgAndRetvalShapes()
1916 const std::vector<InferredShape>& retval_shapes, const Graph& graph, in AssignArgsAndRetvalsToCores() argument
2116 retval_types, retval_shapes); in AssignArgsAndRetvalsToCores()
2224 retval_shapes, graph, replicate_node, flr, in AssignArgsAndRetvalsToCores()
4252 std::vector<InferredShape> arg_shapes, retval_shapes; in RewriteTPUReplicateNode() local
4254 shape_info, *replicate_node, params_info, &arg_shapes, &retval_shapes)); in RewriteTPUReplicateNode()
4264 retval_shapes, *computation, replicate_node, flr, in RewriteTPUReplicateNode()
|