Searched refs:gather_shape (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shape_inference_test.cc | 2426 TF_ASSERT_OK_AND_ASSIGN(Shape gather_shape, in TEST_F()
2436 ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {64, 32}))) in TEST_F()
2437 << ShapeUtil::HumanString(gather_shape); in TEST_F()
2441 TF_ASSERT_OK_AND_ASSIGN(Shape gather_shape, in TEST_F()
2451 ShapeUtil::Equal(gather_shape, ShapeUtil::MakeShape(F32, {32, 48}))) in TEST_F()
2452 << ShapeUtil::HumanString(gather_shape); in TEST_F()
2456 TF_ASSERT_OK_AND_ASSIGN(Shape gather_shape, in TEST_F()
2465 EXPECT_TRUE(ShapeUtil::Equal(gather_shape, in TEST_F()
2467 << ShapeUtil::HumanString(gather_shape); in TEST_F()
2472 Shape gather_shape, in TEST_F()
[all …]
|
| D | hlo_sharding_util.cc | 903 const Shape gather_shape = gather.shape(); in GatherParallelDataOperandSharding() local
908 for (int i = 0, parallel_idx = 0; i < gather_shape.rank(); ++i) { in GatherParallelDataOperandSharding()
1341 const Shape gather_shape = gather.shape(); in GatherParallelOutputDims() local
1343 for (int i = 0, idx_dim = 0; i < gather_shape.dimensions_size(); ++i) { in GatherParallelOutputDims()
|