runtime/onert/test/util/ShapeInference.cc

   1 /*
   2  * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *      http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <gtest/gtest.h>
  18
  19 #include "ir/Layout.h"
  20 #include "util/ShapeInference.h"
  21
  22 using namespace onert::ir;
  23
  24 TEST(ShapeInference, Elementwise)
  25 {
  26   Shape lhs_shape{1, 299, 299, 3};
  27   Shape rhs_shape{3};
  28   auto infered_out_shape = onert::shape_inference::inferEltwiseShape(lhs_shape, rhs_shape);
  29
  30   ASSERT_EQ(infered_out_shape.rank(), 4);
  31   ASSERT_EQ(infered_out_shape.dim(0), 1);
  32   ASSERT_EQ(infered_out_shape.dim(1), 299);
  33   ASSERT_EQ(infered_out_shape.dim(2), 299);
  34   ASSERT_EQ(infered_out_shape.dim(3), 3);
  35 }
  36
  37 TEST(ShapeInference, neg_Elementwise)
  38 {
  39   Shape lhs_shape{1, 299, 299, 3};
  40   Shape rhs_shape{5, 3};
  41   ASSERT_THROW(onert::shape_inference::inferEltwiseShape(lhs_shape, rhs_shape), std::runtime_error);
  42 }
  43
  44 TEST(ShapeInference, Pool2DNodeSame)
  45 {
  46   Shape in_shape{10, 6, 12, 20};
  47   Stride stride{3, 7};
  48   Padding padding{PaddingType::SAME};
  49
  50   operation::Pool2D::Param avg_pool_param{
  51       operation::Pool2D::PoolType::AVG, 3, 6, stride, padding, Activation::NONE};
  52   auto infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, avg_pool_param);
  53
  54   ASSERT_EQ(infered_out_shape.rank(), 4);
  55   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
  56   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
  57   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
  58   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
  59
  60   operation::Pool2D::Param max_pool_param{
  61       operation::Pool2D::PoolType::MAX, 3, 6, stride, padding, Activation::NONE};
  62   infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, max_pool_param);
  63
  64   ASSERT_EQ(infered_out_shape.rank(), 4);
  65   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
  66   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
  67   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
  68   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
  69 }
  70
  71 TEST(ShapeInference, Pool2DNodeValid)
  72 {
  73   Shape in_shape{10, 6, 12, 20};
  74   Stride stride{3, 7};
  75   Padding padding{PaddingType::VALID};
  76
  77   operation::Pool2D::Param avg_pool_param{
  78       operation::Pool2D::PoolType::AVG, 3, 6, stride, padding, Activation::NONE};
  79   auto infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, avg_pool_param);
  80
  81   ASSERT_EQ(infered_out_shape.rank(), 4);
  82   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
  83   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
  84   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
  85   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
  86
  87   operation::Pool2D::Param max_pool_param{
  88       operation::Pool2D::PoolType::MAX, 3, 6, stride, padding, Activation::NONE};
  89   infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, max_pool_param);
  90
  91   ASSERT_EQ(infered_out_shape.rank(), 4);
  92   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
  93   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
  94   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
  95   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
  96 }
  97
  98 TEST(ShapeInference, Pool2DNodeExplicit)
  99 {
 100   Shape in_shape{10, 3, 5, 20};
 101
 102   Stride stride{3, 7};
 103   Padding padding{4, 3, 2, 1};
 104
 105   operation::Pool2D::Param avg_pool_param{
 106       operation::Pool2D::PoolType::AVG, 3, 6, stride, padding, Activation::NONE};
 107   auto infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, avg_pool_param);
 108
 109   ASSERT_EQ(infered_out_shape.rank(), 4);
 110   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 111   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 112   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
 113   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
 114
 115   operation::Pool2D::Param max_pool_param{
 116       operation::Pool2D::PoolType::MAX, 3, 6, stride, padding, Activation::NONE};
 117   infered_out_shape = onert::shape_inference::inferPoolShape(in_shape, max_pool_param);
 118
 119   ASSERT_EQ(infered_out_shape.rank(), 4);
 120   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 121   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 122   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
 123   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 20);
 124 }
 125
 126 TEST(ShapeInference, neg_Pool2DNode_InvalidStride)
 127 {
 128   Shape in_shape{10, 6, 12, 20};
 129   Stride stride{0, 7};
 130   Padding padding{PaddingType::SAME};
 131
 132   operation::Pool2D::Param avg_pool_param{
 133       operation::Pool2D::PoolType::AVG, 3, 6, stride, padding, Activation::NONE};
 134   ASSERT_THROW(onert::shape_inference::inferPoolShape(in_shape, avg_pool_param),
 135                std::runtime_error);
 136 }
 137
 138 TEST(ShapeInference, Conv2D)
 139 {
 140   Shape in_shape{10, 6, 12, 20};
 141   Shape ker_shape{30, 3, 6, 20};
 142
 143   operation::Conv2D::Param param{Stride{3, 7}, Padding{PaddingType::VALID}, Activation::NONE,
 144                                  Dilation{1, 1}};
 145   auto infered_out_shape = onert::shape_inference::inferConv2DShape(in_shape, ker_shape, param);
 146
 147   ASSERT_EQ(infered_out_shape.rank(), 4);
 148   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 149   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 150   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
 151   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 30);
 152
 153   param = operation::Conv2D::Param{Stride{3, 7}, Padding{PaddingType::SAME}, Activation::NONE,
 154                                    Dilation{1, 1}};
 155   infered_out_shape = onert::shape_inference::inferConv2DShape(in_shape, ker_shape, param);
 156
 157   ASSERT_EQ(infered_out_shape.rank(), 4);
 158   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 159   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 160   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
 161   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 30);
 162
 163   param =
 164       operation::Conv2D::Param{Stride{3, 7}, Padding{4, 3, 2, 1}, Activation::NONE, Dilation{1, 1}};
 165   infered_out_shape = onert::shape_inference::inferConv2DShape(in_shape, ker_shape, param);
 166
 167   ASSERT_EQ(infered_out_shape.rank(), 4);
 168   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 169   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 3);
 170   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
 171   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 30);
 172 }
 173
 174 TEST(ShapeInference, neg_Conv2D_InvalidStride)
 175 {
 176   Shape in_shape{10, 6, 12, 20};
 177   Shape ker_shape{30, 3, 6, 20};
 178
 179   operation::Conv2D::Param param{Stride{0, 0}, Padding{PaddingType::VALID}, Activation::NONE,
 180                                  Dilation{1, 1}};
 181   ASSERT_THROW(onert::shape_inference::inferConv2DShape(in_shape, ker_shape, param),
 182                std::runtime_error);
 183 }
 184
 185 TEST(ShapeInference, DepthwiseConv2D)
 186 {
 187   Shape in_shape{10, 6, 12, 20};
 188   Shape ker_shape{1, 3, 6, 60};
 189
 190   operation::DepthwiseConv2D::Param param{Stride{3, 7}, Padding{PaddingType::VALID}, 3,
 191                                           Activation::NONE, Dilation{1, 1}};
 192   auto infered_out_shape =
 193       onert::shape_inference::inferDepthwiseConv2DShape(in_shape, ker_shape, param);
 194
 195   ASSERT_EQ(infered_out_shape.rank(), 4);
 196   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 197   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 198   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 1);
 199   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 60);
 200
 201   param = operation::DepthwiseConv2D::Param{Stride{3, 7}, Padding{PaddingType::SAME}, 3,
 202                                             Activation::NONE, Dilation{1, 1}};
 203   infered_out_shape = onert::shape_inference::inferDepthwiseConv2DShape(in_shape, ker_shape, param);
 204
 205   ASSERT_EQ(infered_out_shape.rank(), 4);
 206   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 207   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 2);
 208   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
 209   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 60);
 210
 211   param = operation::DepthwiseConv2D::Param{Stride{3, 7}, Padding{4, 3, 2, 1}, 3, Activation::NONE,
 212                                             Dilation{1, 1}};
 213   infered_out_shape = onert::shape_inference::inferDepthwiseConv2DShape(in_shape, ker_shape, param);
 214
 215   ASSERT_EQ(infered_out_shape.rank(), 4);
 216   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).N, 10);
 217   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).H, 3);
 218   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).W, 2);
 219   ASSERT_EQ(infered_out_shape.asFeature(Layout::NHWC).C, 60);
 220 }
 221
 222 TEST(ShapeInference, neg_DepthwiseConv2D_InvalidSride)
 223 {
 224   Shape in_shape{10, 6, 12, 20};
 225   Shape ker_shape{1, 3, 6, 60};
 226
 227   operation::DepthwiseConv2D::Param param{Stride{3, 0}, Padding{PaddingType::VALID}, 3,
 228                                           Activation::NONE, Dilation{1, 1}};
 229   ASSERT_THROW(onert::shape_inference::inferDepthwiseConv2DShape(in_shape, ker_shape, param),
 230                std::runtime_error);
 231 }
 232
 233 TEST(ShapeInference, Concat)
 234 {
 235   {
 236     Shape in1{10, 20, 30, 3, 50};
 237     Shape in2{10, 20, 30, 2, 50};
 238     Shape in3{10, 20, 30, 2, 50};
 239
 240     operation::Concat::Param param{3};
 241     auto infered_out_shape = onert::shape_inference::inferConcatShape({in1, in2, in3}, param);
 242
 243     ASSERT_EQ(infered_out_shape.rank(), 5);
 244     ASSERT_EQ(infered_out_shape.dim(0), 10);
 245     ASSERT_EQ(infered_out_shape.dim(1), 20);
 246     ASSERT_EQ(infered_out_shape.dim(2), 30);
 247     ASSERT_EQ(infered_out_shape.dim(3), 7);
 248     ASSERT_EQ(infered_out_shape.dim(4), 50);
 249   }
 250   {
 251     // case 1. when axis < 0
 252     Shape in1{10, 20, 2};
 253     Shape in2{10, 20, 3};
 254
 255     operation::Concat::Param param{-1};
 256     auto infered_out_shape = onert::shape_inference::inferConcatShape({in1, in2}, param);
 257
 258     ASSERT_EQ(infered_out_shape.rank(), 3);
 259     ASSERT_EQ(infered_out_shape.dim(0), 10);
 260     ASSERT_EQ(infered_out_shape.dim(1), 20);
 261     ASSERT_EQ(infered_out_shape.dim(2), 5);
 262   }
 263   {
 264     // case 2. when axis < 0
 265     Shape in1{2, 20, 2};
 266     Shape in2{3, 20, 2};
 267
 268     operation::Concat::Param param{-3};
 269     auto infered_out_shape = onert::shape_inference::inferConcatShape({in1, in2}, param);
 270
 271     ASSERT_EQ(infered_out_shape.rank(), 3);
 272     ASSERT_EQ(infered_out_shape.dim(0), 5);
 273     ASSERT_EQ(infered_out_shape.dim(1), 20);
 274     ASSERT_EQ(infered_out_shape.dim(2), 2);
 275   }
 276 }
 277
 278 TEST(ShapeInference, neg_Concat)
 279 {
 280   {
 281     operation::Concat::Param param{2};
 282     Shape in1{10, 1, 3};
 283     Shape in2{10, 2, 4}; // dim[1] should be 1 but 2
 284
 285     EXPECT_ANY_THROW(onert::shape_inference::inferConcatShape({in1, in2}, param));
 286   }
 287   { // wrong rank
 288     operation::Concat::Param param{2};
 289     Shape in1{10, 2, 3, 4};
 290     Shape in2{10, 2, 4}; // rank should be 4
 291
 292     EXPECT_ANY_THROW(onert::shape_inference::inferConcatShape({in1, in2}, param));
 293   }
 294 }
 295
 296 TEST(ShapeInference, ExpandDims)
 297 {
 298   Shape in_shape{30, 40};
 299
 300   auto check = [&](int32_t axis, Shape &expected) {
 301     auto actual = onert::shape_inference::inferExpandDimsShape(in_shape, axis);
 302
 303     ASSERT_EQ(actual.rank(), 3);
 304     for (int32_t dim = 0; dim < expected.rank(); dim++)
 305       ASSERT_EQ(actual.dim(dim), expected.dim(dim));
 306   };
 307
 308   { // boundary
 309     int32_t axis = 0;
 310     Shape expected{1, 30, 40};
 311     check(axis, expected);
 312   }
 313   { // boundary
 314     int32_t axis = 2;
 315     Shape expected{30, 40, 1};
 316     check(axis, expected);
 317   }
 318   { // inside
 319     int32_t axis = 1;
 320     Shape expected{30, 1, 40};
 321     check(axis, expected);
 322   }
 323   { // negative boundary
 324     int32_t axis = -1;
 325     Shape expected{30, 40, 1};
 326     check(axis, expected);
 327   }
 328   { // negative boundary
 329     int32_t axis = -3;
 330     Shape expected{1, 30, 40};
 331     check(axis, expected);
 332   }
 333 }
 334
 335 TEST(ShapeInference, neg_ExpandDims)
 336 {
 337   Shape in_shape{30, 40};
 338
 339   { // over boundary
 340     int32_t axis = 3;
 341     ASSERT_THROW(onert::shape_inference::inferExpandDimsShape(in_shape, axis), std::runtime_error);
 342   }
 343   { // over boundary
 344     int32_t axis = -4;
 345     ASSERT_THROW(onert::shape_inference::inferExpandDimsShape(in_shape, axis), std::runtime_error);
 346   }
 347 }
 348
 349 TEST(ShapeInference, FullyConnected)
 350 {
 351   Shape in_shape{3, 4, 5, 6};
 352   Shape ker_shape{3, 10};
 353   auto infered_out_shape = onert::shape_inference::inferFullyConnectedShape(in_shape, ker_shape);
 354
 355   ASSERT_EQ(infered_out_shape.rank(), 2);
 356   ASSERT_EQ(infered_out_shape.dim(0), 36);
 357   ASSERT_EQ(infered_out_shape.dim(1), 3);
 358 }
 359
 360 TEST(ShapeInference, Transpose)
 361 {
 362   auto check = [&](Shape &in_shape, std::vector<int> perm, Shape &expected) {
 363     // pre-conditions
 364     ASSERT_EQ(in_shape.rank(), perm.size());
 365     ASSERT_EQ(expected.rank(), perm.size());
 366     auto inferred_out_shape =
 367         onert::shape_inference::inferTransposeShape(in_shape, perm.data(), perm.size());
 368     // post-conditions
 369     ASSERT_EQ(inferred_out_shape.rank(), perm.size());
 370     for (int32_t dim = 0; dim < expected.rank(); dim++)
 371     {
 372       ASSERT_EQ(inferred_out_shape.dim(dim), expected.dim(dim));
 373     }
 374   };
 375   // check for 2-D
 376   {
 377     Shape in_shape{2, 3};
 378     std::vector<int> perm = {1, 0};
 379     Shape expected{3, 2};
 380     // int32_t rank = 2;
 381     check(in_shape, perm, expected);
 382   }
 383   // check for 3-D
 384   {
 385     Shape in_shape{1, 2, 3};
 386     std::vector<int> perm = {2, 0, 1};
 387     Shape expected{3, 1, 2};
 388     // int32_t rank = 3;
 389     check(in_shape, perm, expected);
 390   }
 391   // check for 4-D
 392   {
 393     Shape in_shape{1, 2, 3, 4};
 394     std::vector<int> perm = {1, 3, 0, 2};
 395     Shape expected{2, 4, 1, 3};
 396     // int32_t rank = 4;
 397     check(in_shape, perm, expected);
 398   }
 399 }
 400
 401 TEST(ShapeInference, neg_Transpose)
 402 {
 403   Shape in_shape{1, 2, 3};
 404   // Invalid parameter size
 405   {
 406     std::vector<int> perm = {2, 0, 1, 0};
 407     // int32_t rank = 3;
 408     ASSERT_THROW(onert::shape_inference::inferTransposeShape(in_shape, perm.data(), perm.size()),
 409                  std::runtime_error);
 410   }
 411   // Invalid parameter value
 412   {
 413     std::vector<int> perm = {2, 0, 3};
 414     // int32_t rank = 3;
 415     ASSERT_THROW(onert::shape_inference::inferTransposeShape(in_shape, perm.data(), perm.size()),
 416                  std::runtime_error);
 417   }
 418 }
 419
 420 TEST(ShapeInference, Gather)
 421 {
 422   auto check = [&](Shape &input, Shape &indices, Shape &expected, int32_t axis) {
 423     int rank = input.rank();
 424     auto actual = onert::shape_inference::inferGatherShape(input, indices, axis, rank);
 425
 426     ASSERT_EQ(actual.rank(), expected.rank());
 427
 428     for (int32_t dim = 0; dim < expected.rank(); dim++)
 429       ASSERT_EQ(actual.dim(dim), expected.dim(dim));
 430   };
 431
 432   // check for 2-D, 3-D, axis 0
 433   {
 434     Shape input{3, 4};
 435     Shape indices{1, 1, 2};
 436     int32_t axis = 0;
 437     Shape expected{1, 1, 2, 4};
 438     check(input, indices, expected, axis);
 439   }
 440
 441   // check for 2-D, 3-D, axis 1
 442   {
 443     Shape input{3, 4};
 444     Shape indices{1, 2, 1};
 445     int32_t axis = 1;
 446     Shape expected{3, 1, 2, 1};
 447     check(input, indices, expected, axis);
 448   }
 449
 450   // check for 3-D, 2-D, axis 0
 451   {
 452     Shape input{2, 3, 4};
 453     Shape indices{1, 2};
 454     int32_t axis = 0;
 455     Shape expected{1, 2, 3, 4};
 456     check(input, indices, expected, axis);
 457   }
 458
 459   // check for 3-D, 2-D, axis 2
 460   {
 461     Shape input{2, 3, 4};
 462     Shape indices{2, 1};
 463     int32_t axis = 2;
 464     Shape expected{2, 3, 2, 1};
 465     check(input, indices, expected, axis);
 466   }
 467
 468   // check for 4D, axis 0
 469   {
 470     Shape input{1, 2, 3, 4};
 471     Shape indices{2};
 472     int32_t axis = 0;
 473     Shape expected{2, 2, 3, 4};
 474     check(input, indices, expected, axis);
 475   }
 476 }
 477
 478 TEST(ShapeInference, BCQFullyConnected)
 479 {
 480   auto check = [&](Shape &in_shape, Shape &cluster_shape, std::vector<int> cluster,
 481                    Shape &expected) {
 482     auto actual = onert::shape_inference::inferBCQFullyConnectedShape(in_shape, cluster_shape,
 483                                                                       cluster.data());
 484     ASSERT_EQ(actual.rank(), expected.rank());
 485
 486     for (int32_t dim = 0; dim < expected.rank(); dim++)
 487       ASSERT_EQ(actual.dim(dim), expected.dim(dim));
 488   };
 489
 490   {
 491     Shape in_shape{10, 1};
 492     Shape cluster_shape{3, 2};
 493     std::vector<int> cluster = {1, 10, 2, 10, 3, 10};
 494
 495     Shape expected{30, 1};
 496     check(in_shape, cluster_shape, cluster, expected);
 497   }
 498
 499   {
 500     Shape in_shape{1, 1};
 501     Shape cluster_shape{1, 2};
 502     std::vector<int> cluster = {3, 50};
 503
 504     Shape expected{50, 1};
 505     check(in_shape, cluster_shape, cluster, expected);
 506   }
 507 }
 508
 509 TEST(ShapeInference, BCQGather)
 510 {
 511   auto check = [&](Shape &indices_shape, Shape &cluster_shape, std::vector<int> cluster,
 512                    uint32_t hidden_size, uint32_t axis, int rank, Shape &expected) {
 513     operation::BCQGather::Param param{hidden_size, axis};
 514     auto actual = onert::shape_inference::inferBCQGatherShape(indices_shape, cluster_shape,
 515                                                               cluster.data(), rank, param);
 516     ASSERT_EQ(actual.rank(), expected.rank());
 517
 518     for (int32_t dim = 0; dim < expected.rank(); dim++)
 519       ASSERT_EQ(actual.dim(dim), expected.dim(dim));
 520   };
 521
 522   {
 523     Shape indices_shape{5, 1};
 524     Shape cluster_shape{3, 2};
 525     std::vector<int> cluster = {1, 10, 2, 10, 3, 10};
 526     uint32_t hidden_size = 10;
 527     uint32_t axis = 0;
 528     int rank = 2;
 529
 530     Shape expected{5, 1, 10};
 531     check(indices_shape, cluster_shape, cluster, hidden_size, axis, rank, expected);
 532   }
 533
 534   {
 535     Shape indices_shape{5, 1};
 536     Shape cluster_shape{3, 2};
 537     std::vector<int> cluster = {1, 10, 2, 10, 3, 10};
 538     uint32_t hidden_size = 10;
 539     uint32_t axis = 1;
 540     int rank = 2;
 541
 542     Shape expected{30, 5, 1};
 543     check(indices_shape, cluster_shape, cluster, hidden_size, axis, rank, expected);
 544   }
 545 }