Searched refs:arg_fast_mem (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.h | 313 std::vector<bool>* arg_fast_mem,
359 const std::vector<bool>& arg_fast_mem,
|
| D | distributed_tpu_rewrite_pass.cc | 1919 std::vector<xla::OpSharding>* arg_sharding, std::vector<bool>* arg_fast_mem, in AssignArgsAndRetvalsToCores() argument
1972 arg_fast_mem->resize(args.size()); in AssignArgsAndRetvalsToCores()
2103 (*arg_fast_mem)[i] = is_fast_mem; in AssignArgsAndRetvalsToCores()
2226 arg_fast_mem, retval_sharding, arg_names); in AssignArgsAndRetvalsToCores()
2312 const std::vector<bool>& arg_fast_mem, in BuildCompileNode() argument
2371 arg->set_fast_mem(arg_fast_mem[i]); in BuildCompileNode()
4259 std::vector<bool> arg_fast_mem; in RewriteTPUReplicateNode() local
4265 allow_xla_spmd_partition_, &arg_sharding, &arg_fast_mem, &retval_sharding, in RewriteTPUReplicateNode()
4324 arg_fast_mem, arg_names, retval_sharding, num_cores_per_replica, in RewriteTPUReplicateNode()
|