Searched refs:tuple_sharding (Results 1 – 3 of 3) sorted by relevance
139 HloSharding tuple_sharding = in TEST_F() local143 tuple_sharding.GetAsShapeTree(nested_tuple_shape); in TEST_F()148 EXPECT_IS_OK(tuple_sharding.Validate(nested_tuple_shape, /*num_devices=*/5)); in TEST_F()150 EXPECT_IS_NOT_OK(tuple_sharding.Validate(ShapeUtil::MakeTupleShape({}), in TEST_F()153 EXPECT_IS_NOT_OK(tuple_sharding.Validate(ShapeUtil::MakeShape(F32, {}), in TEST_F()
250 for (auto& tuple_sharding : tuple_elements_) { in UniqueDevice() local251 auto device = tuple_sharding.UniqueDevice(); in UniqueDevice()
792 xla::OpSharding tuple_sharding; in BuildArguments() local793 tuple_sharding.set_type(xla::OpSharding::Type::OpSharding_Type_TUPLE); in BuildArguments()797 *tuple_sharding.add_tuple_shardings() = in BuildArguments()801 tuple_sharding); in BuildArguments()