%178 = "tf.Pack"(%7, %177) {axis = 0 : i64, device = ""} : (tensor<i32>, tensor<i32>) -> tensor<2xi32>
  %179 = "tf.Tile"(%115, %178) {device = ""} : (tensor<?x1xi64>, tensor<2xi32>) -> tensor<?x?xi64>
  %180 = "tf.Mul"(%177, %118) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %181 = "tf.Pack"(%180) {axis = 0 : i64, device = ""} : (tensor<i32>) -> tensor<1xi32>

    int64_t depth = depth_attr.getValues<APInt>()[0].getSExtValue();
    int64_t axis = op.getAxis();
    if (axis == -1) axis = indices_shape.size();

    llvm::SmallVector<int64_t, 4> broadcast_dims(indices_shape.size());
    std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0);
    std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1);

    llvm::SmallVector<int64_t, 4> output_dims =

func.func @expand_dims(%arg0: tensor<2xf32>, %axis: tensor<i32>) -> tensor<1x2xf32> {
  // CHECK: mhlo.reshape
  %0 = "tf.ExpandDims"(%arg0, %axis) : (tensor<2xf32>, tensor<i32>) -> tensor<1x2xf32>
  func.return %0 : tensor<1x2xf32>
}

// -----

// CHECK-LABEL: expand_dims_dynamic
func.func @expand_dims_dynamic(%arg0: tensor<?x?xf32>) -> tensor<?x1x?xf32> {

  %axis = arith.constant dense<2> : tensor<i32>
  %res_ff = "tfl.cumsum"(%arg, %axis) {exclusive = false, reverse = false} : (tensor<1x2x1x3xf32>, tensor<i32>) -> tensor<1x2x1x3xf32>  // Eliminated
  %res_ft = "tfl.cumsum"(%arg, %axis) {exclusive = false, reverse =  true} : (tensor<1x2x1x3xf32>, tensor<i32>) -> tensor<1x2x1x3xf32>  // Eliminated

  func.return %0 : tensor<8xui8>
}

// CHECK-LABEL:   func @torch_index_select(
// CHECK:       %[[AXIS:.+]] = "tf.Const"() <{value = dense<0> : tensor<i64>}> : () -> tensor<i64>
// CHECK:       %[[RES:.+]] = "tf.GatherV2"(%arg0, %arg1, %[[AXIS]]) <{batch_dims = 0 : i64}>
// CHECK:       return %[[RES]]