Searched refs:create_partition_id (Results 1 – 4 of 4) sorted by relevance
115 std::function<HloInstruction*(SpmdBuilder*)> create_partition_id; member
1594 result.create_partition_id = [creator, device_groups](SpmdBuilder* b) { in GetPerGroupCollectiveOpsCreator()1595 return GetInGroupPartitionId(creator.create_partition_id(b), device_groups, in GetPerGroupCollectiveOpsCreator()
667 lhs.state().collective_ops_creator.create_partition_id(&body_b); in PartitionBaseCase()1059 lhs.state().collective_ops_creator.create_partition_id(&body_b); in PartitionBaseCase()
1363 partition_id_(collective_ops_creator_.create_partition_id(&b_)), in SpmdPartitioningVisitor()