Searched defs:arg_sharding (Results 1 – 4 of 4) sorted by relevance
852 const std::optional<xla::HloSharding>& arg_sharding, in XLAShapeForArgument()1053 auto arg_sharding = arg_shardings.find((*input_to_args)[i]); in BuildArguments() local
219 auto arg_sharding = GetXlaShardingFromArg(arg); in IdentifyXlaShardingForComputationInputs() local
2132 std::vector<xla::OpSharding>* arg_sharding, std::vector<bool>* arg_fast_mem, in AssignArgsAndRetvalsToCores()2599 const std::vector<xla::OpSharding>& arg_sharding, in BuildCompileNode()4760 std::vector<xla::OpSharding> arg_sharding; in RewriteTPUReplicateNode() local
2363 for (const auto& arg_sharding : llvm::enumerate(arg_shardings)) { in SetEntryTupleShardings() local