Searched defs:partitioned_input (Results 1 – 6 of 6) sorted by relevance
54 auto partitioned_input = in ReorderReplicateAndPartitionedInputs() local129 getFunction()->walk([](TF::TPUPartitionedInputOp partitioned_input) { in runOnFunction()
92 auto partitioned_input = llvm::dyn_cast_or_null<TF::TPUPartitionedInputOp>( in PartitionResourceReadsWrites() local119 auto partitioned_input = llvm::dyn_cast_or_null<TF::TPUPartitionedInputOp>( in PartitionResourceReadsWrites() local
63 if (auto partitioned_input = in GetXlaShardingFromOperand() local198 if (auto partitioned_input = in GetXlaShardingFromResult() local
370 auto partitioned_input = in HandleFft() local
1671 auto partitioned_input = GetPartitionedHlo(input).PadWithValue( in HandleSort() local1834 auto partitioned_input = GetPartitionedHlo(input).PadWithValue( in HandleCustomCall() local2360 const auto& partitioned_input = in HandleDynamicUpdateSlice() local
257 if (auto partitioned_input = in ExtractInputsForLogicalDevices() local