// CHECK: return %[[einsum]] : tensor<?x2x4x5xf32>

// -----

// dimension_numbers: {
//   offset_dims: 0
//   collapsed_slice_dims: 1
//   start_index_map: 1
// }
func.func @xla_gather(%arg0: tensor<?x2xf32>, %arg1: tensor<1xi32>, %arg2: tensor<2xi32>) -> tensor<*xf32> {

  %16 = "mhlo.gather"(%15, %14) <{dimension_numbers = #mhlo.gather<offset_dims = [0, 1, 2], collapsed_slice_dims = [3], start_index_map = [3], index_vector_dim = 1>, slice_sizes = dense<[4, 8, 64, 1]> : tensor<4xi64>}>...

              builder.getI64Type()),
          start_indices));

  absl::flat_hash_set<int64_t> collapsed_dims;
  collapsed_dims.insert(dimension_numbers.collapsed_slice_dims().begin(),
                        dimension_numbers.collapsed_slice_dims().end());

  // Slice operand by constructed start_indices and slice_sizes.
  auto slice_op = builder.create<TF::SliceOp>(

                              op->slice_sizes, builder)));
    attributes.emplace_back(builder.getNamedAttr(
        "collapsed_slice_dims",
        BuildVhloTensorV1Attr(
            {static_cast<int64_t>(op->collapsed_slice_dims.size())},
            op->collapsed_slice_dims, builder)));
    attributes.emplace_back(builder.getNamedAttr(
        "start_index_map",
        BuildVhloTensorV1Attr(

        %1 = stablehlo.constant dense<2> : tensor<2x3x2xi64>
        %2 = "stablehlo.gather"(%0, %1) {
          dimension_numbers = #stablehlo.gather<
            offset_dims = [2, 3],
            collapsed_slice_dims = [0],
            start_index_map = [1, 0],
            index_vector_dim = 2>,
          slice_sizes = array<i64: 1, 2, 2>,
          indices_are_sorted = false

      return failure();
    }

    // Verify that start_index_map and collapsed_slice_dims contains the same
    // values.
    auto start_index_map = gather_op.getDimensionNumbers().getStartIndexMap();
    auto collapsed_slice_dims =
        gather_op.getDimensionNumbers().getCollapsedSliceDims();
    if (start_index_map.size() != collapsed_slice_dims.size()) {
      return rewriter.notifyMatchFailure(
          gather_op,

func.func @gather(%operand: tensor<3x4x2xi32>, %start_indices: tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>{
  %result = "vhlo.gather_v1"(%operand, %start_indices) <{
    offset_dims = #vhlo.tensor_v1<dense<[2, 3]> : tensor<2xi64>>,
    collapsed_slice_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
    start_index_map = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>,
    index_vector_dim = #vhlo.integer_v1<2 : i64>,

  ) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 3.000000e-01:-5>> {
  %0 = "stablehlo.gather"(%arg0, %arg1) {
    dimension_numbers = #stablehlo.gather<
      offset_dims = [2, 3],
      collapsed_slice_dims = [0, 1],
      start_index_map = [0, 1],
      index_vector_dim = 2>,
    slice_sizes = array<i64: 1, 1, 2, 2>,
    indices_are_sorted = false
  } : (

    ArrayRef<int64_t> slice_sizes = op.getSliceSizes();
    ArrayRef<int64_t> collapsed_slice_dims =
        dim_numbers.getCollapsedSliceDims();
    SmallVector<int64_t> slice_shape;
    for (int64_t i = 0; i < slice_sizes.size(); ++i) {
      // `collapsed_slice_dims` are excluded for slice shape.
      if (!llvm::is_contained(collapsed_slice_dims, i)) {
        slice_shape.push_back(slice_sizes[i]);
      }
    }

// CHECK:         }
func.func @convert_gather(%arg0: tensor<147456xf16>, %arg1: tensor<192x256x1xi32>) -> tensor<192x256xf16> {
  %0 = "mhlo.gather"(%arg0, %arg1) {
    dimension_numbers = #mhlo.gather<
      collapsed_slice_dims = [0],
      index_vector_dim = 2,
			start_index_map = [0],
    >,
    indices_are_sorted = false,
    slice_sizes = dense<1> : tensor<1xi64>