Searched defs:arg_sharding (Results 1 – 5 of 5) sorted by relevance
168 auto arg_sharding = GetXlaShardingFromArg(arg); in IdentifyXlaShardingForComputationInputs() local
835 const absl::optional<xla::HloSharding>& arg_sharding, in XLAShapeForArgument()1026 auto arg_sharding = arg_shardings.find((*input_to_args)[i]); in BuildArguments() local
113 absl::optional<xla::HloSharding> arg_sharding; in GetXlaInputShapes() local
1919 std::vector<xla::OpSharding>* arg_sharding, std::vector<bool>* arg_fast_mem, in AssignArgsAndRetvalsToCores()2311 const std::vector<xla::OpSharding>& arg_sharding, in BuildCompileNode()4258 std::vector<xla::OpSharding> arg_sharding; in RewriteTPUReplicateNode() local
1486 for (auto arg_sharding : llvm::enumerate(arg_shardings)) { in SetEntryTupleShardings() local