Home
last modified time | relevance | path

Searched refs:nested_tuple_shape (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dshape_util_test.cc473 Shape nested_tuple_shape = ShapeUtil::MakeTupleShape( in TEST() local
478 EXPECT_TRUE(ShapeUtil::Equal(nested_tuple_shape, in TEST()
479 ShapeUtil::GetSubshape(nested_tuple_shape, {}))); in TEST()
481 array_shape, ShapeUtil::GetSubshape(nested_tuple_shape, {0}))); in TEST()
484 ShapeUtil::GetSubshape(nested_tuple_shape, {1}))); in TEST()
487 ShapeUtil::GetSubshape(nested_tuple_shape, {2, 0}))); in TEST()
502 Shape nested_tuple_shape = ShapeUtil::MakeTupleShape( in TEST() local
507 EXPECT_FALSE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {})); in TEST()
508 EXPECT_TRUE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {0})); in TEST()
509 EXPECT_FALSE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {1})); in TEST()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dhlo_sharding_test.cc163 Shape nested_tuple_shape = ShapeUtil::MakeTupleShape({ in TEST_F() local
178 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()
H A Dmemory_space_assignment_test.cc3333 Shape nested_tuple_shape = ShapeUtil::MakeTupleShape({shape, tuple_shape}); in TEST_P() local
3350 HloInstruction::CreateParameter(0, nested_tuple_shape, "p")); in TEST_P()
3365 nested_tuple_shape, HloInstruction::FusionKind::kCustom, {p0, p0}, in TEST_P()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dlocal_client_execute_test.cc333 const Shape nested_tuple_shape = in XLA_TEST_F() local
339 auto param = Parameter(&builder, 0, nested_tuple_shape, "param"); in XLA_TEST_F()