Searched refs:tuple_sharding (Results 1 – 6 of 6) sorted by relevance
473 for (const auto& tuple_sharding : tuple_shardings) {476 handle sharding_type = getattr(tuple_sharding, "type");481 dims = getattr(tuple_sharding, "tile_assignment_dimensions")488 devices = getattr(tuple_sharding, "tile_assignment_devices")495 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()
358 for (auto& tuple_sharding : tuple_elements_) { in UniqueDevice() local359 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()
1056 xla::OpSharding tuple_sharding; in BuildArguments() local1057 tuple_sharding.set_type(xla::OpSharding::TUPLE); in BuildArguments()1060 *tuple_sharding.add_tuple_shardings() = in BuildArguments()1074 : tuple_sharding); in BuildArguments()
1797 xla::HloSharding tuple_sharding = xla::HloSharding::Tuple( in TEST_F() local1800 tuple_sharding.ToProto().SerializeAsString()); in TEST_F()