Home
last modified time | relevance | path

Searched refs:tuple_sharding (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_test.cc139 HloSharding tuple_sharding = in TEST_F() local
143 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()
Dhlo_sharding.cc250 for (auto& tuple_sharding : tuple_elements_) { in UniqueDevice() local
251 auto device = tuple_sharding.UniqueDevice(); in UniqueDevice()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler.cc792 xla::OpSharding tuple_sharding; in BuildArguments() local
793 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()