Searched refs:use_spmd (Results 1 – 3 of 3) sorted by relevance
120 StringRef logical_core_0_sharding, bool use_spmd, in IdentifyXlaShardingForComputationInputs() argument158 if (use_spmd) { in IdentifyXlaShardingForComputationInputs()247 StringRef logical_core_0_sharding, bool use_spmd, in IdentifyXlaShardingForComputationOutputs() argument276 if (use_spmd) { in IdentifyXlaShardingForComputationOutputs()315 bool use_spmd = false; in IdentifyXlaShardingForTPUComputation() local318 use_spmd = use_spmd_attr.getValue(); in IdentifyXlaShardingForTPUComputation()320 IdentifyXlaShardingForComputationInputs(logical_core_0_sharding, use_spmd, in IdentifyXlaShardingForTPUComputation()323 IdentifyXlaShardingForComputationOutputs(logical_core_0_sharding, use_spmd, in IdentifyXlaShardingForTPUComputation()
75 bool use_spmd = false; in PartitionResourceReadsWrites() local78 use_spmd = use_spmd_attr.getValue(); in PartitionResourceReadsWrites()80 if (!use_spmd) return; in PartitionResourceReadsWrites()
1974 const bool use_spmd = (UseSpmdForXlaPartitioning(replicate_node) || in AssignArgsAndRetvalsToCores() local2040 if (use_spmd && in AssignArgsAndRetvalsToCores()2166 if (use_spmd) { in AssignArgsAndRetvalsToCores()2211 if (use_spmd && in AssignArgsAndRetvalsToCores()2327 const bool use_spmd = in BuildCompileNode() local2336 proto.set_use_spmd_for_xla_partitioning(use_spmd); in BuildCompileNode()