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

LogicalResult HandleTensorListGatherOp(
    TF::TensorListGatherOp gather,
    const llvm::SmallDenseMap<Value, SizeInfo>& buffer_to_size) {
  auto it = buffer_to_size.find(gather.getInputHandle());
  if (it == buffer_to_size.end()) {
    return gather.emitOpError("unknown tensor list");
  }
  auto buffer = gather.getInputHandle();
  auto result = cutil::GatherElements(gather.getIndices(), buffer,

// CHECK: %[[emb:.*]] = "tfl.gather"(%[[dq_w]], %arg0)
// CHECK: return %[[emb:.*]]

// PerTensor: %[[q_w:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
// PerTensor: %[[dq_w:.*]] = "tfl.dequantize"(%[[q_w]]) : (tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>) -> tensor<64x3x3x3xf32>
// PerTensor: %[[emb:.*]] = "tfl.gather"(%[[dq_w]], %arg0)

      : quant_options_(quant_options), test_mode_(false) {}

  LiftQuantizableSpotsAsFunctionsPass(
      const LiftQuantizableSpotsAsFunctionsPass& other) {
    quant_options_ = other.quant_options_;
    test_mode_ = other.test_mode_;
    op_set_ = other.op_set_;
    initializeForTest();
  }

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in

        if self.same_scale_op == 'concatenate':
          ones = array_ops.ones_like(out)
          out = array_ops.concat([out, ones], 0)
        elif self.same_scale_op == 'gather':
          out = array_ops.gather(out, indices=[0], axis=0)
        elif self.same_scale_op == 'max_pool':
          out = nn_ops.max_pool(out, ksize=3, strides=1, padding='SAME')
        elif self.same_scale_op == 'pad':

// CHECK-NEXT:      %17 = "tfl.gather"(%14, %16) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32>
// CHECK-NEXT:      %18 = "tfl.pseudo_const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64>
// CHECK-NEXT:      %19 = "tfl.gather"(%15, %18) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<1xi64>) -> tensor<1xi32>

      : test_mode_(false), calib_opts_(calib_opts) {}

  InsertCustomAggregationOpsPass(const InsertCustomAggregationOpsPass &other) {
    test_mode_ = other.test_mode_;
    test_case_ = other.test_case_;
    calib_opts_ = other.calib_opts_;
    initializeForTest();
  }

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertCustomAggregationOpsPass)

// designate the dimension indices that were collapsed to produce `tensor_type`.
// The restored dimensions' sizes are 1, according to the semantics of
// `XlaGatherOp (https://www.tensorflow.org/xla/operation_semantics#gather). The
// resulting type's shape has `tensor_type.size() + collapsed_dims.size()`
// dimensions.
RankedTensorType RestoreCollapsedDimensions(
    const RankedTensorType tensor_type,

    // CHECK: %[[GATHER:.*]] = "stablehlo.gather"(%[[CALL]], %[[ARG2]])
    // CHECK-SAME: (tensor<3x4x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>, tensor<2x3x2xi64>) -> tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    // CHECK: %[[DQ:.*]] = "quantfork.dcast"(%[[GATHER]]) : (tensor<2x3x2x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2x2xf32>

// mhlo.dot_general:
//   1. flattened_out_dim = UnsortedSegmentProdOp(operand_shape, out_axes)
//   2. flattened_contracting_dim = UnsortedSegmentProdOp(operand_shape,
//   contracting_axes)
//   3. batch_dimensions = Gather(operand_shape, batch_axes)
//   4. flattened_shape = Concat(batch_dimensions, flattened_out_dim,
//   flattened_contracting_dim)
// The flattened shape for LHS
// is like [batch_dimensions, flattened_out_dimension,