Searched refs:sharding_tree (Results 1 – 2 of 2) sorted by relevance
252 ShapeTree<HloSharding> sharding_tree( in ApplyShardingFromUsers() local277 ? sharding_tree.find({user->tuple_index()}) in ApplyShardingFromUsers()278 : sharding_tree.begin(); in ApplyShardingFromUsers()280 sub_assigned, AssignTreeSharding(&sharding_tree, sharding_tree_begin, in ApplyShardingFromUsers()288 AssignTreeSharding(&sharding_tree, sharding_tree.begin(), in ApplyShardingFromUsers()302 instruction->set_sharding(HloSharding::Tuple(sharding_tree)); in ApplyShardingFromUsers()304 TF_RET_CHECK(sharding_tree.leaf_count() == 1); in ApplyShardingFromUsers()305 instruction->set_sharding(sharding_tree.leaf_begin()->second); in ApplyShardingFromUsers()
302 ShapeTree<HloSharding> sharding_tree(tuple_shape, HloSharding::Replicate()); in TEST_F() local305 for (auto& index_to_sharding : sharding_tree.leaves()) { in TEST_F()308 check(HloSharding::Tuple(sharding_tree)); in TEST_F()