Searched refs:allow_xla_spmd_partition_ (Results 1 – 2 of 2) sorted by relevance
583 static bool allow_xla_spmd_partition_; variable
2328 UseSpmdForXlaPartitioning(replicate_node) && allow_xla_spmd_partition_ && in BuildCompileNode()4265 allow_xla_spmd_partition_, &arg_sharding, &arg_fast_mem, &retval_sharding, in RewriteTPUReplicateNode()4534 bool DistributedTPURewritePass::allow_xla_spmd_partition_ = true; member in tensorflow::DistributedTPURewritePass4548 allow_xla_spmd_partition_ = allow_xla_spmd_partition; in SetDistributedTpuRewritePassOptions()