Searched refs:tuple_sharding (Results 1 – 6 of 6) sorted by relevance
482 for (const auto& tuple_sharding : tuple_shardings) {485 handle sharding_type = getattr(tuple_sharding, "type");490 dims = getattr(tuple_sharding, "tile_assignment_dimensions")497 devices = getattr(tuple_sharding, "tile_assignment_devices")504 getattr(tuple_sharding, "replicate_on_last_tile_dim").cast<bool>());
174 HloSharding tuple_sharding = in TEST_F() local178 tuple_sharding.GetAsShapeTree(nested_tuple_shape); in TEST_F()183 EXPECT_IS_OK(tuple_sharding.Validate(nested_tuple_shape, /*num_devices=*/5)); in TEST_F()185 EXPECT_IS_NOT_OK(tuple_sharding.Validate(ShapeUtil::MakeTupleShape({}), in TEST_F()188 EXPECT_IS_NOT_OK(tuple_sharding.Validate(ShapeUtil::MakeShape(F32, {}), in TEST_F()
434 for (auto& tuple_sharding : tuple_elements_) { in UniqueDevice() local435 auto device = tuple_sharding.UniqueDevice(); in UniqueDevice()
829 for (const HloSharding& tuple_sharding : in TEST_P()831 for (const HloSharding& sub_sharding : tuple_sharding.tuple_elements()) { in TEST_P()
1066 xla::OpSharding tuple_sharding; in BuildArguments() local1067 tuple_sharding.set_type(xla::OpSharding::TUPLE); in BuildArguments()1070 *tuple_sharding.add_tuple_shardings() = in BuildArguments()1084 : tuple_sharding); in BuildArguments()
1864 xla::HloSharding tuple_sharding = xla::HloSharding::Tuple( in TEST_F() local1867 tuple_sharding.ToProto().SerializeAsString()); in TEST_F()